From 0f327e60dc87217a4e63e2762b5577097be3959b Mon Sep 17 00:00:00 2001 From: JJJYmmm <1650675829@qq.com> Date: Mon, 10 Aug 2026 12:42:56 +0800 Subject: [PATCH 01/11] starvla: add shared Qwen-VL conversion and runtime primitives --- .gitattributes | 1 + .gitignore | 3 + .../0001-qwen3vl-vision-parity.patch | 36 + ...002-per-context-native-graph-control.patch | 267 ++ patches/llama.cpp/README.md | 32 + src/models/starvla/normalization.cpp | 171 ++ src/models/starvla/normalization.h | 34 + src/models/starvla/qwen3vl_bridge.cpp | 1667 ++++++++++++ src/models/starvla/qwen3vl_bridge.h | 184 ++ tools/hf2gguf/README.md | 1 + tools/hf2gguf/README_ZH.md | 1 + tools/hf2gguf/environment.yaml | 2 + tools/hf2gguf/starvla/__init__.py | 1 + .../starvla/compare_starvla_actions.py | 82 + .../starvla/convert_starvla_policy_to_gguf.py | 2249 +++++++++++++++++ .../starvla/convert_starvla_qwen_to_gguf.py | 286 +++ tools/hf2gguf/starvla/download_starvla.py | 422 ++++ tools/hf2gguf/starvla/environment.yaml | 14 + .../starvla/inspect_starvla_checkpoint.py | 177 ++ tools/hf2gguf/starvla/starvla_checkpoint.py | 1152 +++++++++ tools/hf2gguf/starvla/starvla_surgery.py | 368 +++ .../starvla/validate_starvla_bundle.py | 1224 +++++++++ tools/llama_cpp/apply_starvla_patches.sh | 115 + 23 files changed, 8489 insertions(+) create mode 100644 .gitattributes create mode 100644 patches/llama.cpp/0001-qwen3vl-vision-parity.patch create mode 100644 patches/llama.cpp/0002-per-context-native-graph-control.patch create mode 100644 patches/llama.cpp/README.md create mode 100644 src/models/starvla/normalization.cpp create mode 100644 src/models/starvla/normalization.h create mode 100644 src/models/starvla/qwen3vl_bridge.cpp create mode 100644 src/models/starvla/qwen3vl_bridge.h create mode 100755 tools/hf2gguf/starvla/__init__.py create mode 100644 tools/hf2gguf/starvla/compare_starvla_actions.py create mode 100755 tools/hf2gguf/starvla/convert_starvla_policy_to_gguf.py create mode 100755 tools/hf2gguf/starvla/convert_starvla_qwen_to_gguf.py create mode 100755 tools/hf2gguf/starvla/download_starvla.py create mode 100644 tools/hf2gguf/starvla/environment.yaml create mode 100755 tools/hf2gguf/starvla/inspect_starvla_checkpoint.py create mode 100755 tools/hf2gguf/starvla/starvla_checkpoint.py create mode 100755 tools/hf2gguf/starvla/starvla_surgery.py create mode 100755 tools/hf2gguf/starvla/validate_starvla_bundle.py create mode 100755 tools/llama_cpp/apply_starvla_patches.sh diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..fe9e1ab --- /dev/null +++ b/.gitattributes @@ -0,0 +1 @@ +patches/**/*.patch -whitespace diff --git a/.gitignore b/.gitignore index ca3edc6..c220675 100644 --- a/.gitignore +++ b/.gitignore @@ -20,6 +20,9 @@ dist/ checkpoint/ checkpoints/ *.safetensors +eval/results/ +/goldens/ +/Testing/ # macOS resource forks ._* diff --git a/patches/llama.cpp/0001-qwen3vl-vision-parity.patch b/patches/llama.cpp/0001-qwen3vl-vision-parity.patch new file mode 100644 index 0000000..a868059 --- /dev/null +++ b/patches/llama.cpp/0001-qwen3vl-vision-parity.patch @@ -0,0 +1,36 @@ +diff --git a/tools/mtmd/models/qwen3vl.cpp b/tools/mtmd/models/qwen3vl.cpp +index fa1100d..5119df6 100644 +--- a/tools/mtmd/models/qwen3vl.cpp ++++ b/tools/mtmd/models/qwen3vl.cpp +@@ -43,8 +43,11 @@ ggml_cgraph * clip_graph_qwen3vl::build() { + cb(inp, "patch_bias", -1); + } + +- // calculate absolute position embedding and apply +- ggml_tensor * learned_pos_embd = resize_position_embeddings(); ++ // Qwen3-VL constructs interpolation coordinates with torch.linspace(0, ++ // num_grid_per_side - 1, size), which is bilinear align_corners=True ++ // without antialiasing. ++ ggml_tensor * learned_pos_embd = resize_position_embeddings( ++ GGML_SCALE_MODE_BILINEAR | GGML_SCALE_FLAG_ALIGN_CORNERS); + learned_pos_embd = ggml_cont_4d( + ctx0, learned_pos_embd, + n_embd * 2, n_patches_x / 2, n_patches_y, batch_size); +@@ -154,7 +157,7 @@ ggml_cgraph * clip_graph_qwen3vl::build() { + layer.deepstack_fc1_w, layer.deepstack_fc1_b, + nullptr, nullptr, + layer.deepstack_fc2_w, layer.deepstack_fc2_b, +- ffn_op_type::FFN_GELU, il); ++ ffn_op_type::FFN_GELU_ERF, il); + + if(!deepstack_features) { + deepstack_features = feat; +@@ -180,7 +183,7 @@ ggml_cgraph * clip_graph_qwen3vl::build() { + model.mm_0_w, model.mm_0_b, + nullptr, nullptr, + model.mm_1_w, model.mm_1_b, +- ffn_op_type::FFN_GELU, -1); ++ ffn_op_type::FFN_GELU_ERF, -1); + + if (deepstack_features) { + embeddings = ggml_concat(ctx0, embeddings, deepstack_features, 0); diff --git a/patches/llama.cpp/0002-per-context-native-graph-control.patch b/patches/llama.cpp/0002-per-context-native-graph-control.patch new file mode 100644 index 0000000..3b3a992 --- /dev/null +++ b/patches/llama.cpp/0002-per-context-native-graph-control.patch @@ -0,0 +1,267 @@ +diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h +index d0c7e5a..d86d0b3 100644 +--- a/ggml/include/ggml-backend.h ++++ b/ggml/include/ggml-backend.h +@@ -215,6 +215,8 @@ extern "C" { + typedef ggml_backend_buffer_type_t * (*ggml_backend_dev_get_extra_bufts_t)(ggml_backend_dev_t device); + // Set the abort callback for the backend + typedef void (*ggml_backend_set_abort_callback_t)(ggml_backend_t backend, ggml_abort_callback abort_callback, void * abort_callback_data); ++ // Enable or disable native graph capture/cache for one backend instance. ++ typedef void (*ggml_backend_set_native_graphs_enabled_t)(ggml_backend_t backend, bool enabled); + // Get a list of feature flags supported by the backend (returns a NULL-terminated array) + struct ggml_backend_feature { + const char * name; +diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh +index 1081750..fb7c279 100644 +--- a/ggml/src/ggml-cuda/common.cuh ++++ b/ggml/src/ggml-cuda/common.cuh +@@ -1373,6 +1373,8 @@ struct ggml_backend_cuda_context { + int curr_stream_no = 0; + + #ifdef USE_CUDA_GRAPH ++ bool cuda_graphs_enabled = true; ++ + // Map from first_node_ptr to cuda_graph - allows multiple graphs per context + // when the computation is split across CPU/GPU (e.g., with --n-cpu-moe) + std::unordered_map> cuda_graphs; +@@ -1405,6 +1407,9 @@ struct ggml_backend_cuda_context { + // Check if any CUDA graph is enabled for this context (used by kernels that need to know + // if graphs are in use without having access to the specific graph key) + bool any_cuda_graph_enabled() const { ++ if (!cuda_graphs_enabled) { ++ return false; ++ } + for (const auto & [key, graph] : cuda_graphs) { + if (graph && graph->is_enabled()) { + return true; +@@ -1415,6 +1420,9 @@ struct ggml_backend_cuda_context { + + // Check if any CUDA graph has an instance for this context + bool any_cuda_graph_has_instance() const { ++ if (!cuda_graphs_enabled) { ++ return false; ++ } + for (const auto & [key, graph] : cuda_graphs) { + if (graph && graph->instance != nullptr) { + return true; +diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu +index 8d21b22..42f7838 100644 +--- a/ggml/src/ggml-cuda/ggml-cuda.cu ++++ b/ggml/src/ggml-cuda/ggml-cuda.cu +@@ -3085,6 +3085,23 @@ static void ggml_backend_cuda_synchronize(ggml_backend_t backend) { + GGML_UNUSED(backend); + } + ++static void ggml_backend_cuda_set_native_graphs_enabled(ggml_backend_t backend, bool enabled) { ++ ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context; ++ ++#ifdef USE_CUDA_GRAPH ++ if (cuda_ctx->cuda_graphs_enabled == enabled) { ++ return; ++ } ++ ++ ggml_backend_cuda_synchronize(backend); ++ cuda_ctx->cuda_graphs.clear(); ++ cuda_ctx->cuda_graphs_enabled = enabled; ++#else ++ GGML_UNUSED(cuda_ctx); ++ GGML_UNUSED(enabled); ++#endif ++} ++ + #ifdef USE_CUDA_GRAPH + static bool ggml_cuda_graph_check_compability(ggml_cgraph * cgraph) { + +@@ -4202,8 +4219,8 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud + } + + #ifdef USE_CUDA_GRAPH +- ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); + if (use_cuda_graph && cuda_graph_update_required) { // End CUDA graph capture ++ ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); + if (graph->graph != nullptr) { + CUDA_CHECK(cudaGraphDestroy(graph->graph)); + graph->graph = nullptr; +@@ -4240,6 +4257,10 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud + + #ifdef USE_CUDA_GRAPH + static bool ggml_cuda_graph_set_enabled(ggml_backend_cuda_context * cuda_ctx, const void * graph_key) { ++ if (!cuda_ctx->cuda_graphs_enabled) { ++ return false; ++ } ++ + ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); + + if (graph->graph == nullptr) { +@@ -4267,10 +4288,8 @@ static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend, + #ifdef USE_CUDA_GRAPH + graph_key = ggml_cuda_graph_get_key(cgraph); + +- ggml_cuda_graph_set_enabled(cuda_ctx, graph_key); +- +- ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); +- if (graph->is_enabled()) { ++ if (ggml_cuda_graph_set_enabled(cuda_ctx, graph_key)) { ++ ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); + const bool graph_compatible = ggml_cuda_graph_check_compability(cgraph); + if (graph_compatible) { + const bool properties_changed = ggml_cuda_graph_update_required(cuda_ctx, cgraph); +@@ -5400,6 +5419,9 @@ static void * ggml_backend_cuda_reg_get_proc_address(ggml_backend_reg_t reg, con + if (strcmp(name, "ggml_backend_get_features") == 0) { + return (void *)ggml_backend_cuda_get_features; + } ++ if (strcmp(name, "ggml_backend_set_native_graphs_enabled") == 0) { ++ return (void *)ggml_backend_cuda_set_native_graphs_enabled; ++ } + return nullptr; + } + +diff --git a/include/llama.h b/include/llama.h +index 2ea2267..aa6c656 100644 +--- a/include/llama.h ++++ b/include/llama.h +@@ -969,6 +969,10 @@ extern "C" { + // Set abort callback + LLAMA_API void llama_set_abort_callback(struct llama_context * ctx, ggml_abort_callback abort_callback, void * abort_callback_data); + ++ // Enable or disable native graph capture/cache for each context backend ++ // that exposes this optional capability. Direct graph computation remains enabled. ++ LLAMA_API void llama_set_backend_native_graphs_enabled(struct llama_context * ctx, bool enabled); ++ + // Wait until all computations are finished + // This is automatically done when using one of the functions below to obtain the computation results + // and is not necessary to call it explicitly in most cases +diff --git a/src/llama-context.cpp b/src/llama-context.cpp +index 71a5939..a705255 100644 +--- a/src/llama-context.cpp ++++ b/src/llama-context.cpp +@@ -1031,6 +1031,21 @@ void llama_context::set_abort_callback(bool (*abort_callback)(void * data), void + } + } + ++void llama_context::set_backend_native_graphs_enabled(bool enabled) { ++ for (auto & backend : backends) { ++ auto * reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend.get())); ++ if (reg == nullptr) { ++ continue; ++ } ++ auto * set_enabled = reinterpret_cast( ++ ggml_backend_reg_get_proc_address( ++ reg, "ggml_backend_set_native_graphs_enabled")); ++ if (set_enabled != nullptr) { ++ set_enabled(backend.get(), enabled); ++ } ++ } ++} ++ + void llama_context::set_embeddings(bool value) { + LLAMA_LOG_DEBUG("%s: value = %d\n", __func__, value); + +@@ -3365,6 +3380,10 @@ void llama_set_abort_callback(llama_context * ctx, bool (*abort_callback)(void * + ctx->set_abort_callback(abort_callback, abort_callback_data); + } + ++void llama_set_backend_native_graphs_enabled(llama_context * ctx, bool enabled) { ++ ctx->set_backend_native_graphs_enabled(enabled); ++} ++ + void llama_set_embeddings(llama_context * ctx, bool embeddings) { + ctx->set_embeddings(embeddings); + } +diff --git a/src/llama-context.h b/src/llama-context.h +index 92d1b0c..d1354a4 100644 +--- a/src/llama-context.h ++++ b/src/llama-context.h +@@ -105,6 +105,7 @@ struct llama_context { + void set_n_threads(int32_t n_threads, int32_t n_threads_batch); + + void set_abort_callback(bool (*abort_callback)(void * data), void * abort_callback_data); ++ void set_backend_native_graphs_enabled(bool enabled); + + void set_embeddings (bool value); + void set_causal_attn(bool value); +diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp +index 513b94f..4e63c93 100644 +--- a/tools/mtmd/clip.cpp ++++ b/tools/mtmd/clip.cpp +@@ -2973,6 +2973,26 @@ void clip_free(clip_ctx * ctx) { + delete ctx; + } + ++void clip_set_backend_native_graphs_enabled(clip_ctx * ctx, bool enabled) { ++ if (ctx == nullptr) { ++ return; ++ } ++ ++ for (ggml_backend_t backend : ctx->backend_ptrs) { ++ ggml_backend_dev_t dev = ggml_backend_get_device(backend); ++ ggml_backend_reg_t reg = dev ? ggml_backend_dev_backend_reg(dev) : nullptr; ++ if (reg == nullptr) { ++ continue; ++ } ++ auto * set_enabled = reinterpret_cast( ++ ggml_backend_reg_get_proc_address( ++ reg, "ggml_backend_set_native_graphs_enabled")); ++ if (set_enabled != nullptr) { ++ set_enabled(backend, enabled); ++ } ++ } ++} ++ + // deprecated + size_t clip_embd_nbytes(const struct clip_ctx * ctx) { + const int32_t nx = ctx->model.hparams.image_size; +diff --git a/tools/mtmd/clip.h b/tools/mtmd/clip.h +index a859b38..d63d43d 100644 +--- a/tools/mtmd/clip.h ++++ b/tools/mtmd/clip.h +@@ -51,6 +51,10 @@ struct clip_init_result clip_init(const char * fname, struct clip_context_params + + void clip_free(struct clip_ctx * ctx); + ++// Enable or disable native graph capture/cache for each CLIP backend that ++// exposes this optional capability. Direct graph computation remains enabled. ++void clip_set_backend_native_graphs_enabled(struct clip_ctx * ctx, bool enabled); ++ + size_t clip_embd_nbytes(const struct clip_ctx * ctx); + size_t clip_embd_nbytes_by_img(const struct clip_ctx * ctx, int img_w, int img_h); + +diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp +index 87da687..5fa1bc3 100644 +--- a/tools/mtmd/mtmd.cpp ++++ b/tools/mtmd/mtmd.cpp +@@ -628,6 +628,19 @@ void mtmd_free(mtmd_context * ctx) { + delete ctx; + } + ++void mtmd_set_backend_native_graphs_enabled(mtmd_context * ctx, bool enabled) { ++ if (ctx == nullptr) { ++ return; ++ } ++ ++ if (ctx->ctx_v != nullptr) { ++ clip_set_backend_native_graphs_enabled(ctx->ctx_v, enabled); ++ } ++ if (ctx->ctx_a != nullptr) { ++ clip_set_backend_native_graphs_enabled(ctx->ctx_a, enabled); ++ } ++} ++ + struct mtmd_tokenizer { + mtmd_context * ctx; + std::vector bitmaps; +diff --git a/tools/mtmd/mtmd.h b/tools/mtmd/mtmd.h +index e364174..7cf3af7 100644 +--- a/tools/mtmd/mtmd.h ++++ b/tools/mtmd/mtmd.h +@@ -110,6 +110,10 @@ MTMD_API mtmd_context * mtmd_init_from_file(const char * mmproj_fname, + + MTMD_API void mtmd_free(mtmd_context * ctx); + ++// Enable or disable native graph capture/cache for each media backend that ++// exposes this optional capability. Direct graph computation remains enabled. ++MTMD_API void mtmd_set_backend_native_graphs_enabled(mtmd_context * ctx, bool enabled); ++ + // whether we need to set non-causal mask before llama_decode + // if chunk is nullptr, we assume the default case where chunk is an image chunk + MTMD_API bool mtmd_decode_use_non_causal(const mtmd_context * ctx, const mtmd_input_chunk * chunk); diff --git a/patches/llama.cpp/README.md b/patches/llama.cpp/README.md new file mode 100644 index 0000000..6601890 --- /dev/null +++ b/patches/llama.cpp/README.md @@ -0,0 +1,32 @@ +# StarVLA llama.cpp patches + +The project pins `third_party/llama.cpp` at commit +`3e941b813b1acbbf06c2203a94ceb33d84748c1e`. The StarVLA runtime needs two +small changes that are not available through that revision's public APIs: + +1. `0001-qwen3vl-vision-parity.patch` matches the upstream Qwen3-VL reference + implementation's position interpolation and exact GELU operations. These + changes are required for action-value parity with the original checkpoint. +2. `0002-per-context-native-graph-control.patch` adds an optional backend API to + disable CUDA graph capture for the text and vision contexts owned by one + StarVLA instance. It prevents retained CUDA graphs from violating long-loop + memory stability gates without globally changing other llama.cpp users. + +Apply both patches before configuring or building the StarVLA runtime: + +```bash +./tools/llama_cpp/apply_starvla_patches.sh +``` + +The command verifies the exact llama.cpp revision and refuses a dirty or +partially patched checkout. It is safe to run again after a complete apply. + +Inspect or remove the overlay with: + +```bash +./tools/llama_cpp/apply_starvla_patches.sh --check +./tools/llama_cpp/apply_starvla_patches.sh --revert +``` + +The parent repository commits only these patch assets. It does not advance or +commit a forked llama.cpp gitlink. diff --git a/src/models/starvla/normalization.cpp b/src/models/starvla/normalization.cpp new file mode 100644 index 0000000..9963a6c --- /dev/null +++ b/src/models/starvla/normalization.cpp @@ -0,0 +1,171 @@ +#include "models/starvla/normalization.h" + +#include +#include +#include + +namespace robotcpp::starvla { + +namespace { + +std::string profile_keys(const NormalizationConfig & config) { + std::ostringstream out; + for (size_t i = 0; i < config.profiles.size(); ++i) { + if (i != 0) { + out << ", "; + } + out << config.profiles[i].key; + } + return out.str(); +} + +} // namespace + +bool validate_normalization_config(const NormalizationConfig & config, int action_dim, std::string & error) { + error.clear(); + if (action_dim <= 0) { + error = "StarVLA normalization requires a positive action dimension"; + return false; + } + if (!std::isfinite(config.binary_threshold) || config.binary_threshold != 0.5f) { + error = "StarVLA normalization binary threshold must use the canonical value 0.5"; + return false; + } + if (config.binary_comparison != "gt" && config.binary_comparison != "ge") { + error = "StarVLA normalization binary comparison must be 'gt' or 'ge'"; + return false; + } + if (config.profiles.empty()) { + error = "StarVLA policy has no normalization profiles"; + return false; + } + + std::vector dimension_kind(static_cast(action_dim), 0); + for (int32_t dim : config.continuous_dimensions) { + if (dim < 0 || dim >= action_dim || dimension_kind[static_cast(dim)] != 0) { + error = "StarVLA continuous action dimensions are invalid or duplicated"; + return false; + } + dimension_kind[static_cast(dim)] = 1; + } + for (int32_t dim : config.binary_dimensions) { + if (dim < 0 || dim >= action_dim || dimension_kind[static_cast(dim)] != 0) { + error = "StarVLA binary action dimensions are invalid or duplicated"; + return false; + } + dimension_kind[static_cast(dim)] = 2; + } + if (std::find(dimension_kind.begin(), dimension_kind.end(), uint8_t{0}) != dimension_kind.end()) { + error = "StarVLA continuous and binary action dimensions must cover every action column"; + return false; + } + + std::vector seen_keys; + seen_keys.reserve(config.profiles.size()); + for (const NormalizationProfile & profile : config.profiles) { + if (profile.key.empty() || std::find(seen_keys.begin(), seen_keys.end(), profile.key) != seen_keys.end()) { + error = "StarVLA normalization profile keys must be non-empty and unique"; + return false; + } + seen_keys.push_back(profile.key); + if (profile.action_q01.size() != static_cast(action_dim) || + profile.action_q99.size() != static_cast(action_dim) || + profile.action_mask.size() != static_cast(action_dim)) { + error = "StarVLA normalization profile shape does not match action dimension: " + profile.key; + return false; + } + for (int dim = 0; dim < action_dim; ++dim) { + const size_t index = static_cast(dim); + if (!std::isfinite(profile.action_q01[index]) || !std::isfinite(profile.action_q99[index])) { + error = "StarVLA normalization quantiles must be finite: " + profile.key; + return false; + } + if (dimension_kind[index] == 1 && profile.action_q99[index] < profile.action_q01[index]) { + error = "StarVLA normalization q99 must not be below q01: " + profile.key; + return false; + } + if (dimension_kind[index] == 1 && profile.action_mask[index] == 0) { + error = "StarVLA continuous action dimension is disabled by the normalization mask: " + profile.key; + return false; + } + if (dimension_kind[index] == 2 && profile.action_mask[index] != 0) { + error = "StarVLA binary action dimension must not use q01/q99 scaling: " + profile.key; + return false; + } + } + } + return true; +} + +const NormalizationProfile * resolve_normalization_profile(const NormalizationConfig & config, + const std::string & profile_key, std::string & error) { + error.clear(); + if (profile_key.empty()) { + if (config.profiles.size() == 1) { + return &config.profiles.front(); + } + error = "StarVLA policy has multiple normalization profiles; select one of: " + profile_keys(config); + return nullptr; + } + for (const NormalizationProfile & profile : config.profiles) { + if (profile.key == profile_key) { + return &profile; + } + } + error = "unknown StarVLA normalization profile '" + profile_key + "'; expected one of: " + profile_keys(config); + return nullptr; +} + +bool denormalize_actions(const NormalizationConfig & config, const std::string & profile_key, + const std::vector & normalized, int horizon, int action_dim, + std::vector & actions, std::string & error) { + actions.clear(); + error.clear(); + if (!validate_normalization_config(config, action_dim, error)) { + return false; + } + if (horizon <= 0 || normalized.size() != static_cast(horizon) * static_cast(action_dim)) { + error = "StarVLA normalized action tensor has an incompatible shape"; + return false; + } + const NormalizationProfile * profile = resolve_normalization_profile(config, profile_key, error); + if (profile == nullptr) { + return false; + } + + std::vector is_binary(static_cast(action_dim), 0); + for (int32_t dim : config.binary_dimensions) { + is_binary[static_cast(dim)] = 1; + } + + actions.resize(normalized.size()); + for (int step = 0; step < horizon; ++step) { + for (int dim = 0; dim < action_dim; ++dim) { + const size_t index = static_cast(step) * static_cast(action_dim) + + static_cast(dim); + const float input_value = normalized[index]; + if (!std::isfinite(input_value)) { + actions.clear(); + error = "StarVLA normalized actions must be finite"; + return false; + } + const float value = + config.clip_actions ? std::clamp(input_value, -1.0f, 1.0f) + : input_value; + if (is_binary[static_cast(dim)] != 0) { + const bool active = + config.binary_comparison == "ge" + ? value >= config.binary_threshold + : value > config.binary_threshold; + actions[index] = active ? 1.0f : 0.0f; + } else { + const float low = profile->action_q01[static_cast(dim)]; + const float high = profile->action_q99[static_cast(dim)]; + actions[index] = (value + 1.0f) * 0.5f * (high - low) + low; + } + } + } + return true; +} + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/normalization.h b/src/models/starvla/normalization.h new file mode 100644 index 0000000..53fe6f1 --- /dev/null +++ b/src/models/starvla/normalization.h @@ -0,0 +1,34 @@ +#pragma once + +#include +#include +#include + +namespace robotcpp::starvla { + +struct NormalizationProfile { + std::string key; + std::vector action_q01; + std::vector action_q99; + std::vector action_mask; +}; + +struct NormalizationConfig { + bool clip_actions = false; + float binary_threshold = 0.5f; + std::string binary_comparison; + std::vector continuous_dimensions; + std::vector binary_dimensions; + std::vector profiles; +}; + +bool validate_normalization_config(const NormalizationConfig & config, int action_dim, std::string & error); + +const NormalizationProfile * resolve_normalization_profile(const NormalizationConfig & config, + const std::string & profile_key, std::string & error); + +bool denormalize_actions(const NormalizationConfig & config, const std::string & profile_key, + const std::vector & normalized, int horizon, int action_dim, + std::vector & actions, std::string & error); + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/qwen3vl_bridge.cpp b/src/models/starvla/qwen3vl_bridge.cpp new file mode 100644 index 0000000..7a3735f --- /dev/null +++ b/src/models/starvla/qwen3vl_bridge.cpp @@ -0,0 +1,1667 @@ +#include "models/starvla/qwen3vl_bridge.h" + +#include "ggml.h" +#include "gguf.h" +#include "llama.h" +#include "llama-model.h" +#include "mtmd.h" +#include "models/starvla/oft_prompt.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace robotcpp::starvla { + +bool qwen_vl_resolve_architecture(const std::string & text_architecture, + const std::string & projector_type, + QwenVLArchitecture & architecture, + std::string & error) { + architecture = QwenVLArchitecture::unknown; + error.clear(); + if (text_architecture == "qwen2vl" && + projector_type == "qwen2.5vl_merger") { + architecture = QwenVLArchitecture::qwen2_5_vl; + return true; + } + if (text_architecture == "qwen3vl" && + projector_type == "qwen3vl_merger") { + architecture = QwenVLArchitecture::qwen3_vl; + return true; + } + if (text_architecture != "qwen2vl" && + text_architecture != "qwen3vl") { + error = "unsupported Qwen-VL text architecture: " + text_architecture; + } else if (projector_type != "qwen2.5vl_merger" && + projector_type != "qwen3vl_merger") { + error = "unsupported Qwen-VL projector type: " + projector_type; + } else { + error = "Qwen-VL text and mmproj architectures do not match"; + } + return false; +} + +const char * qwen_vl_architecture_name(QwenVLArchitecture architecture) { + switch (architecture) { + case QwenVLArchitecture::qwen2_5_vl: + return "qwen2.5-vl"; + case QwenVLArchitecture::qwen3_vl: + return "qwen3-vl"; + case QwenVLArchitecture::unknown: + break; + } + return "unknown"; +} + +bool qwen_vl_is_final_norm_tensor_name(const char * name) noexcept { + return name != nullptr && + (std::strcmp(name, "result_norm") == 0 || + std::strcmp(name, "result_embd_pooled") == 0); +} + +bool qwen_vl_hidden_state_source(QwenVLArchitecture architecture, + int decoder_layer_count, + int deepstack_layer_count, + int32_t hidden_tuple_index, + QwenVLHiddenStateSource & source, + std::string & error) { + source = QwenVLHiddenStateSource{}; + error.clear(); + if (decoder_layer_count <= 0 || hidden_tuple_index <= 0 || + hidden_tuple_index > decoder_layer_count) { + error = "Qwen-VL hidden-state tuple index is out of range"; + return false; + } + if (architecture == QwenVLArchitecture::qwen2_5_vl) { + if (deepstack_layer_count != 0) { + error = "Qwen2.5-VL hidden-state profile cannot contain DeepStack"; + return false; + } + if (hidden_tuple_index == decoder_layer_count) { + source.kind = QwenVLHiddenStateSourceKind::final_norm; + source.layer = -1; + } else { + source.kind = QwenVLHiddenStateSourceKind::decoder_output; + source.layer = hidden_tuple_index - 1; + } + return true; + } + if (architecture == QwenVLArchitecture::qwen3_vl) { + if (deepstack_layer_count <= 0 || + deepstack_layer_count > decoder_layer_count) { + error = "Qwen3-VL DeepStack layer count is incompatible with the model"; + return false; + } + source.kind = hidden_tuple_index <= deepstack_layer_count + ? QwenVLHiddenStateSourceKind::deepstack_output + : QwenVLHiddenStateSourceKind::decoder_output; + source.layer = hidden_tuple_index - 1; + return true; + } + error = "Qwen-VL hidden-state architecture is unknown"; + return false; +} + +bool qwen_vl_select_repetition_penalized_top1( + const float * logits, size_t vocab_size, + const std::vector & full_sequence, float repetition_penalty, + int32_t & token, std::string & error) { + token = -1; + error.clear(); + if (logits == nullptr || vocab_size == 0 || + vocab_size > static_cast(INT32_MAX) || + !std::isfinite(repetition_penalty) || repetition_penalty <= 0.0f) { + error = "Qwen-VL generation selector received an invalid contract"; + return false; + } + + std::vector repeated(vocab_size, uint8_t{0}); + for (int32_t value : full_sequence) { + if (value < 0 || static_cast(value) >= vocab_size) { + error = "Qwen-VL generated sequence contains an out-of-vocabulary token"; + return false; + } + repeated[static_cast(value)] = 1; + } + + float best = -std::numeric_limits::infinity(); + int32_t best_token = -1; + for (size_t index = 0; index < vocab_size; ++index) { + float score = logits[index]; + if (std::isnan(score)) { + error = "Qwen-VL generation logits contain NaN"; + return false; + } + if (repeated[index] != 0) { + score = score < 0.0f ? score * repetition_penalty + : score / repetition_penalty; + } + // torch.argmax returns the first index on ties. + if (best_token < 0 || score > best) { + best = score; + best_token = static_cast(index); + } + } + if (best_token < 0) { + error = "Qwen-VL generation selector did not produce a token"; + return false; + } + token = best_token; + return true; +} + +namespace { + +void quiet_mtmd_log_callback(ggml_log_level level, const char * text, void * user_data) { + (void) user_data; + if (level == GGML_LOG_LEVEL_ERROR) { + std::fputs(text, stderr); + } +} + +struct PreparedMultimodalBatch { + size_t token_count = 0; + llama_pos position_count = 0; + std::vector embeddings; + std::vector positions; + std::vector sequence_counts; + std::vector sequence_values; + std::vector sequences; + std::vector outputs; + std::vector token_ids; + + llama_batch view() { + return { + static_cast(token_count), + nullptr, + embeddings.data(), + positions.data(), + sequence_counts.data(), + sequences.data(), + outputs.data(), + }; + } +}; + +struct BackendPlacement { + bool accelerator_compute = false; + bool cpu_compute = false; +}; + +struct LayerCapture { + BackendPlacement * placement = nullptr; + bool enabled = false; + bool bf16_residual_layer_boundaries = false; + size_t expected_deepstack_layer_count = 0; + size_t token_count = 0; + size_t hidden_size = 0; + std::vector layer_to_slot; + std::vector deepstack_to_slot; + int result_norm_slot = -1; + std::vector values; + std::vector seen; + std::vector rounded_layers; + std::vector rounded_deepstack_layers; + std::string error; + + void disable() { + enabled = false; + token_count = 0; + hidden_size = 0; + layer_to_slot.clear(); + deepstack_to_slot.clear(); + result_norm_slot = -1; + values.clear(); + seen.clear(); + rounded_layers.clear(); + rounded_deepstack_layers.clear(); + error.clear(); + } +}; + +void begin_layer_boundary_tracking(LayerCapture & capture, size_t layer_count) { + if (!capture.bf16_residual_layer_boundaries) { + capture.rounded_layers.clear(); + capture.rounded_deepstack_layers.clear(); + return; + } + capture.rounded_layers.assign(layer_count, uint8_t{0}); + capture.rounded_deepstack_layers.assign( + capture.expected_deepstack_layer_count, uint8_t{0}); +} + +bool validate_layer_boundary_tracking(const LayerCapture & capture, std::string & error) { + if (!capture.bf16_residual_layer_boundaries) { + return true; + } + if (capture.rounded_layers.size() != capture.layer_to_slot.size() || + capture.rounded_deepstack_layers.size() != + capture.expected_deepstack_layer_count || + std::any_of(capture.rounded_layers.begin(), capture.rounded_layers.end(), + [](uint8_t seen) { return seen != 1; }) || + std::any_of(capture.rounded_deepstack_layers.begin(), + capture.rounded_deepstack_layers.end(), + [](uint8_t seen) { return seen != 1; })) { + error = + "Qwen3-VL BF16 residual-boundary roundtrip did not observe " + "every expected l_out/deepstack_out exactly once"; + return false; + } + return true; +} + +bool observe_backend_placement(ggml_tensor * tensor, bool ask, void * user_data) { + if (!ask || tensor == nullptr || tensor->op == GGML_OP_NONE || tensor->buffer == nullptr || + user_data == nullptr) { + return false; + } + auto * placement = static_cast(user_data); + ggml_backend_buffer_type_t buffer_type = ggml_backend_buffer_get_type(tensor->buffer); + ggml_backend_dev_t device = + buffer_type == nullptr ? nullptr : ggml_backend_buft_get_device(buffer_type); + if (device == nullptr) { + return false; + } + switch (ggml_backend_dev_type(device)) { + case GGML_BACKEND_DEVICE_TYPE_GPU: + case GGML_BACKEND_DEVICE_TYPE_IGPU: + case GGML_BACKEND_DEVICE_TYPE_ACCEL: + placement->accelerator_compute = true; + break; + case GGML_BACKEND_DEVICE_TYPE_CPU: + placement->cpu_compute = true; + break; + case GGML_BACKEND_DEVICE_TYPE_META: + break; + } + return false; +} + +int indexed_output_index(const char * name, const char * prefix) { + const size_t prefix_size = std::strlen(prefix); + if (name == nullptr || std::strncmp(name, prefix, prefix_size) != 0) { + return -1; + } + const char * number = name + prefix_size; + if (*number == '\0') { + return -1; + } + errno = 0; + char * end = nullptr; + const long parsed = std::strtol(number, &end, 10); + if (errno != 0 || end == number || *end != '\0' || parsed < 0 || parsed > INT_MAX) { + return -1; + } + return static_cast(parsed); +} + +bool observe_text_and_capture_layers(ggml_tensor * tensor, bool ask, void * user_data) { + auto * capture = static_cast(user_data); + if (capture == nullptr) { + return false; + } + observe_backend_placement(tensor, ask, capture->placement); + if (!capture->enabled || tensor == nullptr) { + return false; + } + + int slot = -1; + const int deepstack_layer = indexed_output_index(tensor->name, "deepstack_out-"); + const int layer = indexed_output_index(tensor->name, "l_out-"); + const bool is_result_norm = + qwen_vl_is_final_norm_tensor_name(tensor->name); + const bool valid_deepstack_layer = + deepstack_layer >= 0 && + static_cast(deepstack_layer) < capture->deepstack_to_slot.size(); + const bool valid_layer = + layer >= 0 && static_cast(layer) < capture->layer_to_slot.size(); + if (valid_deepstack_layer) { + slot = capture->deepstack_to_slot[static_cast(deepstack_layer)]; + } else if (valid_layer) { + slot = capture->layer_to_slot[static_cast(layer)]; + } else if (is_result_norm) { + slot = capture->result_norm_slot; + } + const bool round_layer = + capture->bf16_residual_layer_boundaries && valid_layer; + const bool round_deepstack = + capture->bf16_residual_layer_boundaries && valid_deepstack_layer && + static_cast(deepstack_layer) < + capture->expected_deepstack_layer_count; + if (slot < 0 && !round_layer && !round_deepstack) { + return false; + } + if (ask) { + return true; + } + + try { + if (slot >= 0) { + if (static_cast(slot) >= capture->seen.size()) { + capture->error = "Qwen3-VL layer-capture slot is out of range"; + return false; + } + if (capture->seen[static_cast(slot)] != 0) { + capture->error = + "Qwen3-VL emitted a requested hidden-state output more than once"; + return false; + } + } + if (round_layer && + (static_cast(layer) >= capture->rounded_layers.size() || + capture->rounded_layers[static_cast(layer)] != 0)) { + capture->error = + "Qwen3-VL l_out BF16 roundtrip index is invalid or repeated"; + return false; + } + if (round_deepstack && + (static_cast(deepstack_layer) >= + capture->rounded_deepstack_layers.size() || + capture->rounded_deepstack_layers[static_cast(deepstack_layer)] != 0)) { + capture->error = + "Qwen3-VL deepstack_out BF16 roundtrip index is invalid or repeated"; + return false; + } + if (!ggml_is_contiguous(tensor) || tensor->ne[0] != static_cast(capture->hidden_size) || + tensor->ne[1] != static_cast(capture->token_count) || tensor->ne[2] != 1 || + tensor->ne[3] != 1) { + capture->error = "Qwen3-VL hidden-state output has an incompatible shape or layout"; + return false; + } + if (capture->hidden_size == 0 || + capture->token_count > + std::numeric_limits::max() / capture->hidden_size) { + capture->error = "Qwen3-VL hidden-state capture size overflow"; + return false; + } + const size_t count = capture->token_count * capture->hidden_size; + if (count > std::numeric_limits::max() / sizeof(float) || + (slot >= 0 && + (count == 0 || static_cast(slot) >= capture->values.size() / count))) { + capture->error = "Qwen3-VL hidden-state capture byte range is invalid"; + return false; + } + std::vector rounded(count); + if (tensor->type == GGML_TYPE_F32) { + std::vector source(count); + ggml_backend_tensor_get(tensor, source.data(), 0, count * sizeof(float)); + for (size_t index = 0; index < count; ++index) { + rounded[index] = + ggml_bf16_to_fp32(ggml_fp32_to_bf16(source[index])); + } + } else if (tensor->type == GGML_TYPE_F16) { + if (round_layer || round_deepstack) { + capture->error = + "Qwen3-VL BF16 residual-boundary roundtrip requires F32 tensors"; + return false; + } + std::vector source(count); + ggml_backend_tensor_get(tensor, source.data(), 0, + count * sizeof(ggml_fp16_t)); + for (size_t index = 0; index < count; ++index) { + rounded[index] = ggml_bf16_to_fp32( + ggml_fp32_to_bf16(ggml_fp16_to_fp32(source[index]))); + } + } else if (tensor->type == GGML_TYPE_BF16) { + if (round_layer || round_deepstack) { + capture->error = + "Qwen3-VL BF16 residual-boundary roundtrip requires F32 tensors"; + return false; + } + std::vector source(count); + ggml_backend_tensor_get(tensor, source.data(), 0, + count * sizeof(ggml_bf16_t)); + for (size_t index = 0; index < count; ++index) { + rounded[index] = ggml_bf16_to_fp32(source[index]); + } + } else { + capture->error = std::string("unsupported Qwen3-VL hidden-state output type: ") + + ggml_type_name(tensor->type); + return false; + } + if (round_layer || round_deepstack) { + ggml_backend_tensor_set(tensor, rounded.data(), 0, + count * sizeof(float)); + if (round_layer) { + capture->rounded_layers[static_cast(layer)] = 1; + } else { + capture->rounded_deepstack_layers[ + static_cast(deepstack_layer)] = 1; + } + } + if (slot >= 0) { + float * destination = + capture->values.data() + static_cast(slot) * count; + std::copy(rounded.begin(), rounded.end(), destination); + capture->seen[static_cast(slot)] = 1; + } + return true; + } catch (const std::exception & exception) { + capture->error = std::string("failed to capture Qwen3-VL hidden-state output: ") + + exception.what(); + return false; + } catch (...) { + capture->error = "failed to capture Qwen3-VL hidden-state output"; + return false; + } +} + +int32_t decode_and_synchronize(llama_context * context, llama_batch batch) { + const int32_t result = llama_decode(context, batch); + // A cb_eval capture synchronizes only through its selected node. Drain the + // remaining graph tail before disabling capture, clearing KV state, or + // entering a downstream policy graph. + llama_synchronize(context); + return result; +} + +std::string model_metadata(const llama_model * model, const char * key) { + char value[256] = {}; + const int32_t length = llama_model_meta_val_str(model, key, value, sizeof(value)); + if (length < 0 || static_cast(length) >= sizeof(value)) { + throw std::runtime_error(std::string("missing or oversized Qwen GGUF metadata: ") + key); + } + return value; +} + +std::string gguf_string_metadata(const std::string & path, const char * key) { + gguf_init_params params{}; + params.no_alloc = true; + params.ctx = nullptr; + gguf_context * gguf = gguf_init_from_file(path.c_str(), params); + if (gguf == nullptr) { + throw std::runtime_error("failed to read GGUF metadata: " + path); + } + const int64_t index = gguf_find_key(gguf, key); + if (index < 0 || gguf_get_kv_type(gguf, index) != GGUF_TYPE_STRING) { + gguf_free(gguf); + throw std::runtime_error(std::string("missing GGUF string metadata ") + key + ": " + path); + } + const std::string value = gguf_get_val_str(gguf, index); + gguf_free(gguf); + return value; +} + +std::vector tokenize(const llama_vocab * vocab, const std::string & text, + bool parse_special) { + const int32_t required = -llama_tokenize(vocab, text.data(), static_cast(text.size()), + nullptr, 0, false, parse_special); + if (required <= 0) { + return {}; + } + std::vector tokens(static_cast(required)); + const int32_t count = llama_tokenize(vocab, text.data(), static_cast(text.size()), + tokens.data(), required, false, parse_special); + if (count != required) { + return {}; + } + return tokens; +} + +std::string apply_chat_template(const llama_model * model, + QwenVLArchitecture architecture, + const std::string & content) { + const char * chat_template = llama_model_chat_template(model, nullptr); + if (chat_template == nullptr || chat_template[0] == '\0') { + throw std::runtime_error("Qwen text GGUF has no default chat template"); + } + const llama_chat_message messages[] = { + {"system", "You are a helpful assistant."}, + {"user", content.c_str()}, + }; + const size_t message_offset = + architecture == QwenVLArchitecture::qwen2_5_vl ? 0U : 1U; + const size_t message_count = + architecture == QwenVLArchitecture::qwen2_5_vl ? 2U : 1U; + const int32_t required = + llama_chat_apply_template(chat_template, messages + message_offset, + message_count, true, nullptr, 0); + if (required < 0) { + throw std::runtime_error("failed to size the Qwen chat-template output"); + } + std::vector buffer(static_cast(required) + 1, '\0'); + if (buffer.size() > static_cast(INT32_MAX)) { + throw std::runtime_error("Qwen chat-template output is too large"); + } + const int32_t written = llama_chat_apply_template( + chat_template, messages + message_offset, message_count, true, + buffer.data(), static_cast(buffer.size())); + if (written != required) { + throw std::runtime_error("failed to apply the Qwen chat template"); + } + return std::string(buffer.data(), static_cast(written)); +} + +struct PackedImageLayout { + size_t row_bytes = 0; + size_t stride_bytes = 0; + size_t packed_bytes = 0; +}; + +bool validate_image(const Qwen3VLImageView & image, const Qwen3VLBridgeConfig & config, + PackedImageLayout & layout, std::string & error) { + (void) config; + layout = PackedImageLayout{}; + if (image.data == nullptr || image.channels != 3 || image.width <= 0 || + image.height <= 0) { + error = "Qwen3-VL bridge requires a non-empty RGB image"; + return false; + } + const size_t width = static_cast(image.width); + const size_t height = static_cast(image.height); + if (width > std::numeric_limits::max() / 3U) { + error = "Qwen3-VL image row size overflow"; + return false; + } + layout.row_bytes = width * 3U; + if (image.stride_bytes < 0 || + (image.stride_bytes > 0 && + static_cast(image.stride_bytes) < layout.row_bytes)) { + error = "Qwen3-VL image stride is smaller than a packed RGB row"; + return false; + } + layout.stride_bytes = image.stride_bytes > 0 + ? static_cast(image.stride_bytes) + : layout.row_bytes; + if (height > std::numeric_limits::max() / layout.row_bytes || + (height > 1U && + height - 1U > + (std::numeric_limits::max() - layout.row_bytes) / + layout.stride_bytes)) { + error = "Qwen3-VL image buffer size overflow"; + return false; + } + layout.packed_bytes = height * layout.row_bytes; + return true; +} + +std::vector pack_image(const Qwen3VLImageView & image, + const PackedImageLayout & layout) { + std::vector packed(layout.packed_bytes); + for (int row = 0; row < image.height; ++row) { + std::memcpy(packed.data() + static_cast(row) * layout.row_bytes, + image.data + static_cast(row) * layout.stride_bytes, + layout.row_bytes); + } + return packed; +} + +void tokenize_multimodal_prompt(const Qwen3VLBridgeConfig & config, + QwenVLArchitecture architecture, + const llama_model * model, mtmd_context * vision, + const std::vector & images, + const std::string & instruction, + mtmd::input_chunks & chunks) { + std::vector> packed_images; + packed_images.reserve(images.size()); + mtmd::bitmaps bitmaps; + for (const Qwen3VLImageView & image : images) { + std::string validation_error; + PackedImageLayout layout; + if (!validate_image(image, config, layout, validation_error)) { + throw std::runtime_error(validation_error); + } + packed_images.push_back(pack_image(image, layout)); + bitmaps.entries.emplace_back(static_cast(image.width), + static_cast(image.height), + packed_images.back().data()); + if (bitmaps.entries.back().ptr == nullptr) { + throw std::runtime_error("failed to create a Qwen3-VL image bitmap"); + } + } + + const std::string content = + build_qwen_media_content(images.size(), instruction, mtmd_default_marker()); + const std::string formatted = + apply_chat_template(model, architecture, content); + mtmd_input_text input_text{}; + input_text.text = formatted.c_str(); + input_text.add_special = false; + input_text.parse_special = true; + chunks.ptr.reset(mtmd_input_chunks_init()); + if (chunks.ptr == nullptr) { + throw std::runtime_error("failed to allocate Qwen3-VL multimodal input chunks"); + } + std::vector bitmap_ptrs = bitmaps.c_ptr(); + const int32_t tokenize_result = mtmd_tokenize( + vision, chunks.ptr.get(), &input_text, bitmap_ptrs.data(), bitmap_ptrs.size()); + if (tokenize_result != 0) { + throw std::runtime_error("failed to tokenize the Qwen3-VL multimodal prompt"); + } +} + +const char * compiled_backend_name() { +#if defined(GGML_USE_CUDA) + return "cuda"; +#elif defined(GGML_USE_METAL) + return "metal"; +#else + return "cpu"; +#endif +} + +} // namespace + +struct Qwen3VLBridge::Impl { + Qwen3VLBridgeConfig config; + QwenVLArchitecture architecture = QwenVLArchitecture::unknown; + size_t deepstack_layer_count = 0; + llama_model * model = nullptr; + llama_context * context = nullptr; + mtmd_context * vision = nullptr; + const llama_vocab * vocab = nullptr; + bool backend_initialized = false; + BackendPlacement text_placement; + BackendPlacement vision_placement; + LayerCapture layer_capture; + mutable std::string backend_name = "unknown"; + + void refresh_backend_name() const { + const bool accelerator = + text_placement.accelerator_compute && vision_placement.accelerator_compute; + const bool cpu = text_placement.cpu_compute || vision_placement.cpu_compute; + if (accelerator && !cpu) { + backend_name = compiled_backend_name(); + } else if (!text_placement.accelerator_compute && + !vision_placement.accelerator_compute && + text_placement.cpu_compute && vision_placement.cpu_compute) { + backend_name = "cpu"; + } else if (text_placement.accelerator_compute || + vision_placement.accelerator_compute) { + backend_name = "mixed"; + } else { + backend_name = "unknown"; + } + } + + ~Impl() { + if (vision != nullptr) { + mtmd_free(vision); + vision = nullptr; + } + if (context != nullptr) { + llama_free(context); + context = nullptr; + } + if (model != nullptr) { + llama_model_free(model); + model = nullptr; + } + if (backend_initialized) { + llama_backend_free(); + backend_initialized = false; + } + } +}; + +namespace { + +void copy_token_embedding(const ggml_tensor * token_embeddings, llama_token token, + size_t hidden_size, float * destination) { + if (token_embeddings == nullptr || destination == nullptr || token < 0 || + token_embeddings->ne[0] != static_cast(hidden_size) || + token >= token_embeddings->ne[1] || !ggml_is_contiguous(token_embeddings) || + token_embeddings->buffer == nullptr) { + throw std::runtime_error("Qwen3-VL token embedding table is incompatible"); + } + const size_t row_stride = token_embeddings->nb[1]; + if (row_stride == 0 || + static_cast(token) > std::numeric_limits::max() / row_stride) { + throw std::runtime_error("Qwen3-VL token embedding row offset overflow"); + } + const size_t row_offset = static_cast(token) * row_stride; + switch (token_embeddings->type) { + case GGML_TYPE_F32: + ggml_backend_tensor_get(token_embeddings, destination, row_offset, + hidden_size * sizeof(float)); + break; + case GGML_TYPE_F16: { + std::vector row(hidden_size); + ggml_backend_tensor_get(token_embeddings, row.data(), row_offset, + hidden_size * sizeof(ggml_fp16_t)); + for (size_t index = 0; index < hidden_size; ++index) { + destination[index] = ggml_fp16_to_fp32(row[index]); + } + break; + } + case GGML_TYPE_BF16: { + std::vector row(hidden_size); + ggml_backend_tensor_get(token_embeddings, row.data(), row_offset, + hidden_size * sizeof(ggml_bf16_t)); + for (size_t index = 0; index < hidden_size; ++index) { + destination[index] = ggml_bf16_to_fp32(row[index]); + } + break; + } + default: + throw std::runtime_error(std::string("unsupported Qwen3-VL token embedding type: ") + + ggml_type_name(token_embeddings->type)); + } +} + +PreparedMultimodalBatch prepare_multimodal_batch(const Qwen3VLBridgeConfig & config, + llama_model * model, + mtmd_context * vision, + const mtmd::input_chunks & chunks) { + if (model == nullptr || vision == nullptr || !mtmd_decode_use_mrope(vision)) { + throw std::runtime_error("Qwen3-VL single-batch decode requires M-RoPE components"); + } + + PreparedMultimodalBatch prepared; + for (size_t chunk_index = 0; chunk_index < chunks.size(); ++chunk_index) { + const mtmd_input_chunk * chunk = chunks[chunk_index]; + const size_t chunk_tokens = mtmd_input_chunk_get_n_tokens(chunk); + if (chunk_tokens == 0 || chunk_tokens > std::numeric_limits::max() - + prepared.token_count) { + throw std::runtime_error("Qwen3-VL multimodal chunk has an invalid token count"); + } + prepared.token_count += chunk_tokens; + } + if (prepared.token_count == 0 || prepared.token_count > static_cast(INT32_MAX)) { + throw std::runtime_error("Qwen3-VL multimodal prompt token count is invalid"); + } + + const size_t hidden_size = static_cast(config.hidden_size); + const size_t input_size = static_cast(config.input_embedding_size); + if (hidden_size == 0 || input_size != static_cast(llama_model_n_embd_inp(model)) || + input_size < hidden_size || + prepared.token_count > std::numeric_limits::max() / input_size) { + throw std::runtime_error("Qwen3-VL input embedding dimensions are incompatible"); + } + if (prepared.token_count > std::numeric_limits::max() / 4U) { + throw std::runtime_error("Qwen3-VL M-RoPE position buffer size overflow"); + } + + prepared.embeddings.assign(prepared.token_count * input_size, 0.0f); + prepared.positions.resize(prepared.token_count * 4U); + prepared.sequence_counts.assign(prepared.token_count, 1); + prepared.sequence_values.assign(prepared.token_count, 0); + prepared.sequences.resize(prepared.token_count); + prepared.outputs.assign(prepared.token_count, int8_t{0}); + prepared.token_ids.assign(prepared.token_count, static_cast(-1)); + for (size_t index = 0; index < prepared.token_count; ++index) { + prepared.sequences[index] = &prepared.sequence_values[index]; + } + + const ggml_tensor * token_embeddings = model->get_tensor("token_embd.weight"); + size_t token_offset = 0; + llama_pos position_offset = 0; + for (size_t chunk_index = 0; chunk_index < chunks.size(); ++chunk_index) { + const mtmd_input_chunk * chunk = chunks[chunk_index]; + const size_t chunk_tokens = mtmd_input_chunk_get_n_tokens(chunk); + const llama_pos chunk_positions = mtmd_input_chunk_get_n_pos(chunk); + if (chunk_positions <= 0 || + position_offset > std::numeric_limits::max() - chunk_positions) { + throw std::runtime_error("Qwen3-VL multimodal positions overflow"); + } + + const mtmd_input_chunk_type type = mtmd_input_chunk_get_type(chunk); + if (type == MTMD_INPUT_CHUNK_TYPE_TEXT) { + size_t text_token_count = 0; + const llama_token * tokens = + mtmd_input_chunk_get_tokens_text(chunk, &text_token_count); + if (tokens == nullptr || text_token_count != chunk_tokens || + chunk_positions != static_cast(chunk_tokens)) { + throw std::runtime_error("Qwen3-VL text chunk contract is incompatible"); + } + for (size_t local_index = 0; local_index < chunk_tokens; ++local_index) { + const size_t global_index = token_offset + local_index; + copy_token_embedding(token_embeddings, tokens[local_index], hidden_size, + prepared.embeddings.data() + global_index * input_size); + prepared.token_ids[global_index] = tokens[local_index]; + const llama_pos position = + position_offset + static_cast(local_index); + for (size_t axis = 0; axis < 3U; ++axis) { + prepared.positions[axis * prepared.token_count + global_index] = position; + } + prepared.positions[prepared.token_count * 3U + global_index] = 0; + } + } else if (type == MTMD_INPUT_CHUNK_TYPE_IMAGE) { + const mtmd_image_tokens * image_tokens = + mtmd_input_chunk_get_tokens_image(chunk); + if (image_tokens == nullptr || + mtmd_image_tokens_get_n_tokens(image_tokens) != chunk_tokens) { + throw std::runtime_error("Qwen3-VL image chunk contract is incompatible"); + } + if (mtmd_encode_chunk(vision, chunk) != 0) { + throw std::runtime_error("failed to encode a Qwen3-VL image chunk"); + } + const float * image_embeddings = mtmd_get_output_embd(vision); + if (image_embeddings == nullptr) { + throw std::runtime_error("Qwen3-VL image encoder returned no embeddings"); + } + const size_t image_element_count = chunk_tokens * input_size; + float * destination = + prepared.embeddings.data() + token_offset * input_size; + for (size_t element = 0; element < image_element_count; ++element) { + destination[element] = ggml_bf16_to_fp32( + ggml_fp32_to_bf16(image_embeddings[element])); + } + for (size_t local_index = 0; local_index < chunk_tokens; ++local_index) { + const size_t global_index = token_offset + local_index; + const mtmd_decoder_pos position = mtmd_image_tokens_get_decoder_pos( + image_tokens, position_offset, local_index); + prepared.positions[global_index] = static_cast(position.t); + prepared.positions[prepared.token_count + global_index] = + static_cast(position.y); + prepared.positions[prepared.token_count * 2U + global_index] = + static_cast(position.x); + prepared.positions[prepared.token_count * 3U + global_index] = + static_cast(position.z); + } + } else { + throw std::runtime_error("Qwen3-VL prompt contains an unsupported media chunk"); + } + token_offset += chunk_tokens; + position_offset += chunk_positions; + } + if (token_offset != prepared.token_count) { + throw std::runtime_error("Qwen3-VL prepared batch token count mismatch"); + } + prepared.position_count = position_offset; + return prepared; +} + +void export_prepared_inputs(const Qwen3VLBridgeConfig & config, + const llama_vocab * vocab, + const mtmd::input_chunks & chunks, + const PreparedMultimodalBatch & prepared, + std::vector & input_ids, + std::vector & attention_mask, + std::vector & image_grid_thw) { + const std::vector image_pad_tokens = + tokenize(vocab, "<|image_pad|>", true); + if (image_pad_tokens.size() != 1) { + throw std::runtime_error( + "Qwen3-VL vocabulary does not expose a unique <|image_pad|> token"); + } + if (config.image_spatial_merge_size <= 0) { + throw std::runtime_error("Qwen3-VL image spatial merge size is invalid"); + } + + input_ids.reserve(prepared.token_ids.size()); + for (llama_token token : prepared.token_ids) { + input_ids.push_back(token < 0 ? static_cast(image_pad_tokens.front()) + : static_cast(token)); + } + attention_mask.assign(prepared.token_count, uint8_t{1}); + + image_grid_thw.reserve(static_cast(config.expected_image_count) * 3U); + size_t observed_images = 0; + for (size_t chunk_index = 0; chunk_index < chunks.size(); ++chunk_index) { + const mtmd_input_chunk * chunk = chunks[chunk_index]; + if (mtmd_input_chunk_get_type(chunk) != MTMD_INPUT_CHUNK_TYPE_IMAGE) { + continue; + } + const mtmd_image_tokens * image_tokens = + mtmd_input_chunk_get_tokens_image(chunk); + if (image_tokens == nullptr) { + throw std::runtime_error("Qwen3-VL image chunk has no token grid"); + } + const size_t image_token_count = + mtmd_image_tokens_get_n_tokens(image_tokens); + uint32_t max_x = 0; + uint32_t max_y = 0; + for (size_t token = 0; token < image_token_count; ++token) { + const mtmd_decoder_pos position = + mtmd_image_tokens_get_decoder_pos(image_tokens, 0, token); + if (position.t != 0 || position.z != 0) { + throw std::runtime_error( + "Qwen3-VL image token grid does not use the expected M-RoPE layout"); + } + max_x = std::max(max_x, position.x); + max_y = std::max(max_y, position.y); + } + const size_t merged_width = static_cast(max_x) + 1U; + const size_t merged_height = static_cast(max_y) + 1U; + const size_t merge = static_cast(config.image_spatial_merge_size); + if (merged_width == 0 || merged_height == 0 || + merged_width > static_cast(INT64_MAX) / merge || + merged_height > static_cast(INT64_MAX) / merge || + merged_width > std::numeric_limits::max() / merged_height || + merged_width * merged_height != image_token_count) { + throw std::runtime_error("Qwen3-VL image token grid is incompatible"); + } + image_grid_thw.push_back(1); + image_grid_thw.push_back(static_cast(merged_height * merge)); + image_grid_thw.push_back(static_cast(merged_width * merge)); + ++observed_images; + } + if (observed_images != static_cast(config.expected_image_count)) { + throw std::runtime_error("Qwen3-VL image grid count does not match the policy"); + } +} + +} // namespace + +Qwen3VLBridge::Qwen3VLBridge(std::unique_ptr impl) : impl_(std::move(impl)) {} + +Qwen3VLBridge::~Qwen3VLBridge() = default; + +std::unique_ptr Qwen3VLBridge::load(const Qwen3VLBridgeConfig & config, + std::string & error) { + error.clear(); + const bool action_config_valid = + config.action_token.empty() ? config.action_token_id == -1 : config.action_token_id >= 0; + if (config.text_path.empty() || config.mmproj_path.empty() || config.bundle_uuid.empty() || + config.hidden_size <= 0 || config.input_embedding_size <= 0 || config.vocab_size <= 0 || + !action_config_valid || + config.expected_image_count <= 0 || config.image_min_tokens <= 0 || + config.image_max_tokens < config.image_min_tokens || + config.image_spatial_merge_size <= 0 || config.n_ctx <= 0 || + config.n_batch <= 0) { + error = "Qwen3-VL bridge configuration is incomplete"; + return nullptr; + } + + std::unique_ptr impl(new Impl()); + impl->config = config; + try { + const std::string mmproj_uuid = + gguf_string_metadata(config.mmproj_path, "general.source.uuid"); + if (mmproj_uuid != config.bundle_uuid) { + throw std::runtime_error("Qwen3-VL mmproj bundle UUID does not match the policy"); + } + const std::string projector_type = + gguf_string_metadata(config.mmproj_path, "clip.projector_type"); + + llama_backend_init(); + impl->backend_initialized = true; + llama_model_params model_params = llama_model_default_params(); + model_params.n_gpu_layers = -1; + impl->model = llama_model_load_from_file(config.text_path.c_str(), model_params); + if (impl->model == nullptr) { + throw std::runtime_error("failed to load Qwen3-VL text GGUF: " + config.text_path); + } + if (model_metadata(impl->model, "general.source.uuid") != config.bundle_uuid) { + throw std::runtime_error("Qwen3-VL text bundle UUID does not match the policy"); + } + std::string profile_error; + if (!qwen_vl_resolve_architecture( + model_metadata(impl->model, "general.architecture"), + projector_type, impl->architecture, profile_error)) { + throw std::runtime_error(profile_error); + } + if (llama_model_n_embd_out(impl->model) != config.hidden_size || + llama_model_n_embd_inp(impl->model) != config.input_embedding_size) { + throw std::runtime_error("Qwen3-VL text embedding dimensions do not match the policy"); + } + if (config.input_embedding_size % config.hidden_size != 0) { + throw std::runtime_error( + "Qwen-VL input embedding width is not an integral hidden-state layout"); + } + const int deepstack_layer_count = + config.input_embedding_size / config.hidden_size - 1; + if ((impl->architecture == QwenVLArchitecture::qwen2_5_vl && + deepstack_layer_count != 0) || + (impl->architecture == QwenVLArchitecture::qwen3_vl && + (deepstack_layer_count <= 0 || + deepstack_layer_count > llama_model_n_layer(impl->model)))) { + throw std::runtime_error( + "Qwen-VL input embedding layout does not match the detected architecture"); + } + impl->deepstack_layer_count = + static_cast(deepstack_layer_count); + for (int layer = 0; layer < llama_model_n_layer(impl->model); ++layer) { + ggml_backend_dev_t device = impl->model->dev_layer(layer); + if (device == nullptr) { + throw std::runtime_error("Qwen3-VL text layer has no assigned backend device"); + } + const enum ggml_backend_dev_type type = ggml_backend_dev_type(device); + if (type == GGML_BACKEND_DEVICE_TYPE_CPU) { + impl->text_placement.cpu_compute = true; + } else if (type == GGML_BACKEND_DEVICE_TYPE_GPU || + type == GGML_BACKEND_DEVICE_TYPE_IGPU || + type == GGML_BACKEND_DEVICE_TYPE_ACCEL) { + impl->text_placement.accelerator_compute = true; + } + } + if (ggml_backend_dev_t output_device = impl->model->dev_output()) { + const enum ggml_backend_dev_type type = ggml_backend_dev_type(output_device); + if (type == GGML_BACKEND_DEVICE_TYPE_CPU) { + impl->text_placement.cpu_compute = true; + } else if (type == GGML_BACKEND_DEVICE_TYPE_GPU || + type == GGML_BACKEND_DEVICE_TYPE_IGPU || + type == GGML_BACKEND_DEVICE_TYPE_ACCEL) { + impl->text_placement.accelerator_compute = true; + } + } + + impl->vocab = llama_model_get_vocab(impl->model); + if (impl->vocab == nullptr) { + throw std::runtime_error("Qwen3-VL text GGUF has no vocabulary"); + } + if (llama_vocab_n_tokens(impl->vocab) != config.vocab_size) { + throw std::runtime_error("Qwen3-VL text vocabulary size does not match the policy"); + } + if (!config.action_token.empty()) { + const std::vector action_tokens = + tokenize(impl->vocab, config.action_token, true); + if (action_tokens.size() != 1 || action_tokens.front() != config.action_token_id) { + throw std::runtime_error( + "Qwen3-VL action token mapping does not match the policy metadata"); + } + } + + llama_context_params context_params = llama_context_default_params(); + context_params.n_ctx = static_cast(config.n_ctx); + context_params.n_batch = static_cast(config.n_batch); + // Layer capture expects one complete l_out/deepstack_out tensor per decode call. + context_params.n_ubatch = static_cast(config.n_batch); + context_params.n_threads = config.n_threads; + context_params.n_threads_batch = config.n_threads; + context_params.pooling_type = LLAMA_POOLING_TYPE_NONE; + context_params.embeddings = false; + // Match the official Qwen3-VL BF16 inference cache instead of llama's F16 default. + context_params.type_k = GGML_TYPE_BF16; + context_params.type_v = GGML_TYPE_BF16; + context_params.flash_attn_type = config.flash_text_attention + ? LLAMA_FLASH_ATTN_TYPE_ENABLED + : LLAMA_FLASH_ATTN_TYPE_DISABLED; + impl->layer_capture.placement = &impl->text_placement; + impl->layer_capture.bf16_residual_layer_boundaries = + config.bf16_residual_layer_boundaries; + if (config.bf16_residual_layer_boundaries) { + impl->layer_capture.expected_deepstack_layer_count = + impl->deepstack_layer_count; + } + context_params.cb_eval = observe_text_and_capture_layers; + context_params.cb_eval_user_data = &impl->layer_capture; + impl->context = llama_init_from_model(impl->model, context_params); + if (impl->context == nullptr) { + throw std::runtime_error("failed to create Qwen3-VL text context"); + } + if (config.disable_text_backend_native_graphs) { + llama_set_backend_native_graphs_enabled(impl->context, false); + } + if (llama_n_batch(impl->context) != static_cast(config.n_batch) || + llama_n_ubatch(impl->context) != static_cast(config.n_batch)) { + throw std::runtime_error( + "Qwen3-VL text context did not preserve the requested batch/ubatch contract"); + } + + mtmd_context_params vision_params = mtmd_context_params_default(); + vision_params.use_gpu = true; + vision_params.print_timings = config.verbosity >= 1; + vision_params.n_threads = config.n_threads; + vision_params.image_min_tokens = config.image_min_tokens; + vision_params.image_max_tokens = config.image_max_tokens; + vision_params.cb_eval = observe_backend_placement; + vision_params.cb_eval_user_data = &impl->vision_placement; + mtmd_log_set(config.verbosity >= 1 ? nullptr : quiet_mtmd_log_callback, nullptr); + impl->vision = mtmd_init_from_file(config.mmproj_path.c_str(), impl->model, vision_params); + if (impl->vision == nullptr) { + throw std::runtime_error("failed to load Qwen3-VL mmproj GGUF: " + config.mmproj_path); + } + if (config.disable_vision_backend_native_graphs) { + mtmd_set_backend_native_graphs_enabled(impl->vision, false); + } + impl->refresh_backend_name(); + + if (config.verbosity >= 1) { + std::fprintf(stderr, + "%s: architecture=%s backend=%s hidden=%d input_embd=%d " + "deepstack=%zu images=%d image_tokens=%d..%d " + "n_ctx=%u n_batch=%u n_ubatch=%u kv=bf16 " + "text_native_graph_disable_requested=%s " + "vision_native_graph_disable_requested=%s\n", + __func__, qwen_vl_architecture_name(impl->architecture), + impl->backend_name.c_str(), + llama_model_n_embd_out(impl->model), + llama_model_n_embd_inp(impl->model), + impl->deepstack_layer_count, config.expected_image_count, + config.image_min_tokens, config.image_max_tokens, + llama_n_ctx(impl->context), llama_n_batch(impl->context), + llama_n_ubatch(impl->context), + config.disable_text_backend_native_graphs ? "true" : "false", + config.disable_vision_backend_native_graphs ? "true" : "false"); + } + } catch (const std::exception & exception) { + error = exception.what(); + return nullptr; + } + return std::unique_ptr(new Qwen3VLBridge(std::move(impl))); +} + +bool Qwen3VLBridge::extract_token_embeddings(const std::vector & images, + const std::string & instruction, int32_t token_id, + size_t token_count, std::vector & embeddings, + std::string & error) { + embeddings.clear(); + error.clear(); + if (impl_ == nullptr || impl_->model == nullptr || impl_->context == nullptr || + impl_->vision == nullptr || impl_->vocab == nullptr) { + error = "Qwen3-VL bridge is not initialized"; + return false; + } + if (impl_->config.action_token.empty() || impl_->config.action_token_id < 0) { + error = "Qwen3-VL bridge was configured without an action token"; + return false; + } + if (images.size() != static_cast(impl_->config.expected_image_count)) { + error = "Qwen3-VL image count does not match the policy"; + return false; + } + if (token_count == 0) { + error = "Qwen3-VL requested token embedding count must be positive"; + return false; + } + if (token_id != impl_->config.action_token_id) { + error = "Qwen3-VL requested token ID does not match the policy"; + return false; + } + + try { + mtmd::input_chunks chunks; + tokenize_multimodal_prompt(impl_->config, impl_->architecture, + impl_->model, impl_->vision, images, + instruction, chunks); + PreparedMultimodalBatch prepared = + prepare_multimodal_batch(impl_->config, impl_->model, impl_->vision, chunks); + + std::vector matches; + for (size_t index = 0; index < prepared.token_ids.size(); ++index) { + if (prepared.token_ids[index] == token_id) { + matches.push_back(index); + } + } + if (matches.size() < token_count) { + throw std::runtime_error("Qwen3-VL prompt contains fewer target tokens than requested"); + } + matches.erase(matches.begin(), matches.end() - static_cast(token_count)); + if (prepared.token_count > llama_n_batch(impl_->context)) { + throw std::runtime_error( + "Qwen3-VL multimodal prompt exceeds n_batch; increase --n-batch for single-batch decode"); + } + if (prepared.token_count > static_cast(llama_n_ctx(impl_->context)) || + prepared.position_count > static_cast(llama_n_ctx(impl_->context))) { + throw std::runtime_error("Qwen3-VL multimodal prompt exceeds n_ctx"); + } + + const int layer_count = llama_model_n_layer(impl_->model); + if (layer_count <= 0) { + throw std::runtime_error("Qwen3-VL model has no decoder layers"); + } + LayerCapture & capture = impl_->layer_capture; + capture.enabled = false; + capture.token_count = prepared.token_count; + capture.hidden_size = static_cast(impl_->config.hidden_size); + capture.layer_to_slot.assign(static_cast(layer_count), -1); + capture.deepstack_to_slot.assign(static_cast(layer_count), -1); + capture.result_norm_slot = -1; + if (impl_->architecture == QwenVLArchitecture::qwen2_5_vl) { + capture.result_norm_slot = 0; + } else { + capture.layer_to_slot.back() = 0; + } + capture.values.assign(prepared.token_count * capture.hidden_size, 0.0f); + capture.seen.assign(1, uint8_t{0}); + begin_layer_boundary_tracking(capture, static_cast(layer_count)); + capture.error.clear(); + capture.enabled = true; + std::fill(prepared.outputs.begin(), prepared.outputs.end(), int8_t{1}); + + llama_memory_clear(llama_get_memory(impl_->context), true); + llama_set_embeddings(impl_->context, true); + llama_batch batch = prepared.view(); + const int32_t decode_result = decode_and_synchronize(impl_->context, batch); + capture.enabled = false; + if (decode_result != 0) { + throw std::runtime_error("failed to evaluate the Qwen3-VL multimodal batch"); + } + if (!capture.error.empty()) { + throw std::runtime_error(capture.error); + } + std::string boundary_error; + if (!validate_layer_boundary_tracking(capture, boundary_error)) { + throw std::runtime_error(boundary_error); + } + if (capture.seen.size() != 1 || capture.seen.front() == 0) { + throw std::runtime_error( + "Qwen-VL did not expose the final conditioning output"); + } + + embeddings.resize(token_count * capture.hidden_size); + for (size_t output_index = 0; output_index < matches.size(); ++output_index) { + const float * hidden = + capture.values.data() + matches[output_index] * capture.hidden_size; + std::copy_n(hidden, capture.hidden_size, + embeddings.data() + output_index * capture.hidden_size); + } + impl_->layer_capture.disable(); + llama_set_embeddings(impl_->context, false); + impl_->refresh_backend_name(); + return true; + } catch (const std::exception & exception) { + llama_synchronize(impl_->context); + impl_->layer_capture.disable(); + llama_set_embeddings(impl_->context, false); + llama_memory_clear(llama_get_memory(impl_->context), true); + embeddings.clear(); + error = exception.what(); + return false; + } +} + +bool Qwen3VLBridge::extract_full_hidden_states( + const std::vector & images, const std::string & instruction, + std::vector & hidden_states, std::vector & attention_mask, + std::string & error) { + hidden_states.clear(); + attention_mask.clear(); + error.clear(); + if (impl_ == nullptr || impl_->model == nullptr || impl_->context == nullptr || + impl_->vision == nullptr || impl_->vocab == nullptr) { + error = "Qwen3-VL bridge is not initialized"; + return false; + } + if (images.size() != static_cast(impl_->config.expected_image_count)) { + error = "Qwen3-VL image count does not match the policy"; + return false; + } + + try { + mtmd::input_chunks chunks; + tokenize_multimodal_prompt(impl_->config, impl_->architecture, + impl_->model, impl_->vision, images, + instruction, chunks); + PreparedMultimodalBatch prepared = + prepare_multimodal_batch(impl_->config, impl_->model, impl_->vision, chunks); + if (prepared.token_count > llama_n_batch(impl_->context)) { + throw std::runtime_error( + "Qwen3-VL multimodal prompt exceeds n_batch; increase --n-batch for single-batch decode"); + } + if (prepared.token_count > static_cast(llama_n_ctx(impl_->context)) || + prepared.position_count > static_cast(llama_n_ctx(impl_->context))) { + throw std::runtime_error("Qwen3-VL multimodal prompt exceeds n_ctx"); + } + const size_t hidden_size = static_cast(impl_->config.hidden_size); + if (prepared.token_count > std::numeric_limits::max() / hidden_size) { + throw std::runtime_error("Qwen3-VL hidden-state buffer size overflow"); + } + + const int layer_count = llama_model_n_layer(impl_->model); + if (layer_count <= 0) { + throw std::runtime_error("Qwen3-VL model has no decoder layers"); + } + LayerCapture & capture = impl_->layer_capture; + capture.enabled = false; + capture.token_count = prepared.token_count; + capture.hidden_size = hidden_size; + capture.layer_to_slot.assign(static_cast(layer_count), -1); + capture.deepstack_to_slot.assign(static_cast(layer_count), -1); + capture.result_norm_slot = -1; + if (impl_->architecture == QwenVLArchitecture::qwen2_5_vl) { + capture.result_norm_slot = 0; + } else { + capture.layer_to_slot.back() = 0; + } + capture.values.assign(prepared.token_count * hidden_size, 0.0f); + capture.seen.assign(1, uint8_t{0}); + begin_layer_boundary_tracking(capture, static_cast(layer_count)); + capture.error.clear(); + capture.enabled = true; + std::fill(prepared.outputs.begin(), prepared.outputs.end(), int8_t{1}); + + llama_memory_clear(llama_get_memory(impl_->context), true); + llama_set_embeddings(impl_->context, true); + llama_batch batch = prepared.view(); + const int32_t decode_result = decode_and_synchronize(impl_->context, batch); + capture.enabled = false; + if (decode_result != 0) { + throw std::runtime_error("failed to evaluate the Qwen3-VL multimodal batch"); + } + if (!capture.error.empty()) { + throw std::runtime_error(capture.error); + } + std::string boundary_error; + if (!validate_layer_boundary_tracking(capture, boundary_error)) { + throw std::runtime_error(boundary_error); + } + if (capture.seen.size() != 1 || capture.seen.front() == 0) { + throw std::runtime_error( + "Qwen-VL did not expose the final conditioning output"); + } + + hidden_states = std::move(capture.values); + attention_mask.assign(prepared.token_count, uint8_t{1}); + impl_->layer_capture.disable(); + llama_set_embeddings(impl_->context, false); + impl_->refresh_backend_name(); + return true; + } catch (const std::exception & exception) { + llama_synchronize(impl_->context); + impl_->layer_capture.disable(); + llama_set_embeddings(impl_->context, false); + llama_memory_clear(llama_get_memory(impl_->context), true); + hidden_states.clear(); + attention_mask.clear(); + error = exception.what(); + return false; + } +} + +bool Qwen3VLBridge::extract_layer_hidden_states( + const std::vector & images, const std::string & instruction, + const std::vector & hidden_tuple_indices, std::vector & hidden_states, + std::vector & attention_mask, std::string & error) { + hidden_states.clear(); + attention_mask.clear(); + error.clear(); + if (impl_ == nullptr || impl_->model == nullptr || impl_->context == nullptr || + impl_->vision == nullptr || impl_->vocab == nullptr) { + error = "Qwen3-VL bridge is not initialized"; + return false; + } + if (images.size() != static_cast(impl_->config.expected_image_count)) { + error = "Qwen3-VL image count does not match the policy"; + return false; + } + + const int model_layer_count = llama_model_n_layer(impl_->model); + if (hidden_tuple_indices.empty() || model_layer_count <= 0 || + hidden_tuple_indices.size() > static_cast(model_layer_count)) { + error = "Qwen3-VL requested hidden-state layer set is incompatible with the model"; + return false; + } + std::vector layer_to_slot(static_cast(model_layer_count), -1); + std::vector deepstack_to_slot(static_cast(model_layer_count), -1); + if (impl_->deepstack_layer_count > + static_cast(std::numeric_limits::max())) { + error = "Qwen-VL DeepStack layer count exceeds the supported range"; + return false; + } + const int deepstack_layer_count = + static_cast(impl_->deepstack_layer_count); + int result_norm_slot = -1; + for (size_t slot = 0; slot < hidden_tuple_indices.size(); ++slot) { + const int32_t tuple_index = hidden_tuple_indices[slot]; + QwenVLHiddenStateSource source; + if (!qwen_vl_hidden_state_source( + impl_->architecture, model_layer_count, deepstack_layer_count, + tuple_index, source, error)) { + return false; + } + if (source.kind == QwenVLHiddenStateSourceKind::final_norm) { + if (result_norm_slot >= 0) { + error = "Qwen-VL hidden-state tuple indices must be unique"; + return false; + } + result_norm_slot = static_cast(slot); + continue; + } + if (source.layer < 0 || source.layer >= model_layer_count) { + error = "Qwen-VL hidden-state source layer is out of range"; + return false; + } + std::vector & target = + source.kind == QwenVLHiddenStateSourceKind::deepstack_output + ? deepstack_to_slot + : layer_to_slot; + if (target[static_cast(source.layer)] >= 0) { + error = "Qwen-VL hidden-state tuple indices must be unique"; + return false; + } + target[static_cast(source.layer)] = static_cast(slot); + } + + try { + mtmd::input_chunks chunks; + tokenize_multimodal_prompt(impl_->config, impl_->architecture, + impl_->model, impl_->vision, images, + instruction, chunks); + PreparedMultimodalBatch prepared = + prepare_multimodal_batch(impl_->config, impl_->model, impl_->vision, chunks); + if (prepared.token_count > llama_n_batch(impl_->context)) { + throw std::runtime_error( + "Qwen3-VL multimodal prompt exceeds n_batch; increase --n-batch for single-batch decode"); + } + if (prepared.token_count > static_cast(llama_n_ctx(impl_->context)) || + prepared.position_count > static_cast(llama_n_ctx(impl_->context))) { + throw std::runtime_error("Qwen3-VL multimodal prompt exceeds n_ctx"); + } + + const size_t hidden_size = static_cast(impl_->config.hidden_size); + const size_t requested_layers = hidden_tuple_indices.size(); + if (prepared.token_count > std::numeric_limits::max() / hidden_size || + prepared.token_count * hidden_size > + std::numeric_limits::max() / requested_layers) { + throw std::runtime_error("Qwen3-VL layerwise hidden-state buffer size overflow"); + } + LayerCapture & capture = impl_->layer_capture; + capture.enabled = false; + capture.token_count = prepared.token_count; + capture.hidden_size = hidden_size; + capture.layer_to_slot = layer_to_slot; + capture.deepstack_to_slot = deepstack_to_slot; + capture.result_norm_slot = result_norm_slot; + capture.values.assign(requested_layers * prepared.token_count * hidden_size, 0.0f); + capture.seen.assign(requested_layers, uint8_t{0}); + begin_layer_boundary_tracking(capture, + static_cast(model_layer_count)); + capture.error.clear(); + capture.enabled = true; + std::fill(prepared.outputs.begin(), prepared.outputs.end(), int8_t{1}); + + llama_memory_clear(llama_get_memory(impl_->context), true); + llama_set_embeddings(impl_->context, true); + llama_batch batch = prepared.view(); + const int32_t decode_result = decode_and_synchronize(impl_->context, batch); + capture.enabled = false; + if (decode_result != 0) { + throw std::runtime_error("failed to evaluate the Qwen3-VL multimodal batch"); + } + if (!capture.error.empty()) { + throw std::runtime_error(capture.error); + } + std::string boundary_error; + if (!validate_layer_boundary_tracking(capture, boundary_error)) { + throw std::runtime_error(boundary_error); + } + if (std::any_of(capture.seen.begin(), capture.seen.end(), + [](uint8_t seen) { return seen == 0; })) { + throw std::runtime_error( + "Qwen3-VL did not expose every requested hidden-state output"); + } + + hidden_states = std::move(capture.values); + attention_mask.assign(prepared.token_count, uint8_t{1}); + impl_->layer_capture.disable(); + llama_set_embeddings(impl_->context, false); + impl_->refresh_backend_name(); + return true; + } catch (const std::exception & exception) { + llama_synchronize(impl_->context); + impl_->layer_capture.disable(); + llama_set_embeddings(impl_->context, false); + llama_memory_clear(llama_get_memory(impl_->context), true); + hidden_states.clear(); + attention_mask.clear(); + error = exception.what(); + return false; + } +} + +bool Qwen3VLBridge::generate_autoregressive( + const std::vector & images, + const std::string & instruction, + const QwenVLGenerationConfig & generation, + QwenVLGenerationResult & result, std::string & error) { + result = QwenVLGenerationResult{}; + error.clear(); + if (impl_ == nullptr || impl_->model == nullptr || impl_->context == nullptr || + impl_->vision == nullptr || impl_->vocab == nullptr) { + error = "Qwen-VL bridge is not initialized"; + return false; + } + if (images.size() != static_cast(impl_->config.expected_image_count)) { + error = "Qwen-VL image count does not match the policy"; + return false; + } + if (generation.max_length == 0 || + generation.max_length > static_cast(llama_n_ctx(impl_->context)) || + generation.max_length > static_cast(INT32_MAX) || + generation.top_k != 1 || generation.eos_token_ids.empty() || + !std::isfinite(generation.repetition_penalty) || + generation.repetition_penalty <= 0.0f) { + error = "Qwen-VL autoregressive generation configuration is incompatible"; + return false; + } + std::vector eos_seen(static_cast(impl_->config.vocab_size), + uint8_t{0}); + for (int32_t eos : generation.eos_token_ids) { + if (eos < 0 || eos >= impl_->config.vocab_size || + eos_seen[static_cast(eos)] != 0) { + error = "Qwen-VL generation EOS token set is invalid"; + return false; + } + eos_seen[static_cast(eos)] = 1; + } + + try { + mtmd::input_chunks chunks; + tokenize_multimodal_prompt(impl_->config, impl_->architecture, + impl_->model, impl_->vision, images, + instruction, chunks); + PreparedMultimodalBatch prepared = + prepare_multimodal_batch(impl_->config, impl_->model, impl_->vision, + chunks); + if (prepared.token_count > llama_n_batch(impl_->context)) { + throw std::runtime_error( + "Qwen-VL multimodal prompt exceeds n_batch; increase --n-batch"); + } + if (prepared.token_count > generation.max_length || + prepared.token_count > static_cast(llama_n_ctx(impl_->context)) || + prepared.position_count > + static_cast(llama_n_ctx(impl_->context))) { + throw std::runtime_error( + "Qwen-VL multimodal prompt exceeds the FAST max_length/n_ctx contract"); + } + + std::vector input_ids; + std::vector attention_mask; + std::vector image_grid_thw; + export_prepared_inputs(impl_->config, impl_->vocab, chunks, prepared, + input_ids, attention_mask, image_grid_thw); + if (input_ids.size() != prepared.token_count) { + throw std::runtime_error( + "Qwen-VL multimodal prompt token export is inconsistent"); + } + result.prompt_token_count = prepared.token_count; + result.full_sequence.reserve(generation.max_length); + for (int64_t input_id : input_ids) { + if (input_id < 0 || input_id >= impl_->config.vocab_size) { + throw std::runtime_error( + "Qwen-VL multimodal prompt contains an out-of-vocabulary token"); + } + result.full_sequence.push_back(static_cast(input_id)); + } + if (prepared.token_count == generation.max_length) { + impl_->refresh_backend_name(); + return true; + } + + impl_->layer_capture.disable(); + llama_set_embeddings(impl_->context, false); + llama_memory_clear(llama_get_memory(impl_->context), true); + std::fill(prepared.outputs.begin(), prepared.outputs.end(), int8_t{0}); + prepared.outputs.back() = 1; + llama_batch prompt_batch = prepared.view(); + if (decode_and_synchronize(impl_->context, prompt_batch) != 0) { + throw std::runtime_error( + "failed to evaluate the Qwen-VL autoregressive prompt"); + } + + while (result.full_sequence.size() < generation.max_length) { + const float * logits = llama_get_logits_ith(impl_->context, -1); + int32_t next = -1; + std::string selection_error; + if (!qwen_vl_select_repetition_penalized_top1( + logits, static_cast(impl_->config.vocab_size), + result.full_sequence, generation.repetition_penalty, next, + selection_error)) { + throw std::runtime_error(selection_error); + } + result.full_sequence.push_back(next); + result.continuation.push_back(next); + if (eos_seen[static_cast(next)] != 0 || + result.full_sequence.size() == generation.max_length) { + break; + } + + const size_t generation_index = result.continuation.size() - 1U; + if (generation_index > + static_cast(std::numeric_limits::max() - + prepared.position_count)) { + throw std::runtime_error( + "Qwen-VL autoregressive M-RoPE position overflow"); + } + llama_token token = static_cast(next); + llama_pos position = + prepared.position_count + static_cast(generation_index); + int32_t sequence_count = 1; + llama_seq_id sequence_value = 0; + llama_seq_id * sequence = &sequence_value; + int8_t output = 1; + llama_batch token_batch{ + 1, + &token, + nullptr, + &position, + &sequence_count, + &sequence, + &output, + }; + if (decode_and_synchronize(impl_->context, token_batch) != 0) { + throw std::runtime_error( + "failed to evaluate an incremental Qwen-VL generation token"); + } + } + + llama_memory_clear(llama_get_memory(impl_->context), true); + impl_->refresh_backend_name(); + return true; + } catch (const std::exception & exception) { + llama_synchronize(impl_->context); + impl_->layer_capture.disable(); + llama_set_embeddings(impl_->context, false); + llama_memory_clear(llama_get_memory(impl_->context), true); + result = QwenVLGenerationResult{}; + error = exception.what(); + return false; + } +} + +void Qwen3VLBridge::reset() { + if (impl_ != nullptr && impl_->context != nullptr) { + llama_synchronize(impl_->context); + impl_->layer_capture.disable(); + llama_set_embeddings(impl_->context, false); + llama_memory_clear(llama_get_memory(impl_->context), true); + } +} + +const char * Qwen3VLBridge::backend_name() const { + if (impl_ == nullptr) { + return "unknown"; + } + impl_->refresh_backend_name(); + return impl_->backend_name.c_str(); +} + +const char * Qwen3VLBridge::text_attention_mode_name() const { + return impl_ != nullptr && impl_->config.flash_text_attention ? "flash" : "non_flash"; +} + +QwenVLArchitecture Qwen3VLBridge::architecture() const { + return impl_ != nullptr ? impl_->architecture : QwenVLArchitecture::unknown; +} + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/qwen3vl_bridge.h b/src/models/starvla/qwen3vl_bridge.h new file mode 100644 index 0000000..ec6f3eb --- /dev/null +++ b/src/models/starvla/qwen3vl_bridge.h @@ -0,0 +1,184 @@ +#pragma once + +#include +#include +#include +#include +#include + +namespace robotcpp::starvla { + +enum class QwenVLArchitecture { + unknown, + qwen2_5_vl, + qwen3_vl, +}; + +enum class QwenVLHiddenStateSourceKind { + decoder_output, + deepstack_output, + final_norm, +}; + +struct QwenVLHiddenStateSource { + QwenVLHiddenStateSourceKind kind = + QwenVLHiddenStateSourceKind::decoder_output; + int layer = -1; +}; + +// Resolve the paired llama.cpp text and mtmd projector profiles. StarVLA +// supports Qwen2.5-VL and Qwen3-VL only; mismatched text/mmproj files fail +// before either component is evaluated. +bool qwen_vl_resolve_architecture(const std::string & text_architecture, + const std::string & projector_type, + QwenVLArchitecture & architecture, + std::string & error); + +const char * qwen_vl_architecture_name(QwenVLArchitecture architecture); + +// llama.cpp names the final normalized decoder state `result_norm`, then +// renames the same tensor when embedding output with pooling type NONE is +// enabled. Both names identify the Qwen2.5 hidden_states[-1] boundary. +bool qwen_vl_is_final_norm_tensor_name(const char * name) noexcept; + +// Map one Transformers 4.57 hidden_states tuple index to the corresponding +// llama.cpp graph output. Index zero (the embedding input) is intentionally not +// exposed. Qwen2.5-VL has no DeepStack and its final tuple item is result_norm; +// Qwen3-VL retains the recorder/alias behavior used by the existing parity +// contract. +bool qwen_vl_hidden_state_source(QwenVLArchitecture architecture, + int decoder_layer_count, + int deepstack_layer_count, + int32_t hidden_tuple_index, + QwenVLHiddenStateSource & source, + std::string & error); + +struct Qwen3VLImageView { + const uint8_t * data = nullptr; + int width = 0; + int height = 0; + int channels = 0; + int stride_bytes = 0; +}; + +struct Qwen3VLBridgeConfig { + std::string text_path; + std::string mmproj_path; + std::string bundle_uuid; + int hidden_size = 0; + int input_embedding_size = 0; + int vocab_size = 0; + std::string action_token; + int32_t action_token_id = -1; + int expected_image_count = 0; + int image_min_tokens = 0; + int image_max_tokens = 0; + int image_spatial_merge_size = 0; + int n_ctx = 2048; + int n_batch = 2048; + int n_threads = 0; + int verbosity = 0; + // OFT uses plain flash attention to meet its final-action parity profile. + // Other StarVLA variants retain non-flash text attention. + bool flash_text_attention = false; + // Round each F32 decoder residual output, plus DeepStack outputs when the + // detected architecture has them, through BF16 RNE before it feeds the next + // layer. Intra-layer computation keeps llama.cpp's backend-default profile. + bool bf16_residual_layer_boundaries = false; + // Repeated text decode can rebuild graphs with transient node keys. Disable + // native graph capture/cache for the text context so those keys cannot + // accumulate backend graph instances. Direct graph computation continues. + bool disable_text_backend_native_graphs = false; + // The vision encoder rebuilds its graph for every image. Disable native + // graph capture/cache when its transient graph keys are not stable. + bool disable_vision_backend_native_graphs = false; +}; + +struct QwenVLGenerationConfig { + size_t max_length = 0; + std::vector eos_token_ids; + int top_k = 0; + float repetition_penalty = 0.0f; +}; + +struct QwenVLGenerationResult { + size_t prompt_token_count = 0; + std::vector full_sequence; + std::vector continuation; +}; + +// Implements the deterministic token choice used by the official FAST +// generation profile: Hugging Face repetition penalty over the full sequence, +// followed by top_k=1. Exposed so the generation contract can be tested +// without loading a multi-gigabyte Qwen checkpoint. +bool qwen_vl_select_repetition_penalized_top1( + const float * logits, size_t vocab_size, + const std::vector & full_sequence, float repetition_penalty, + int32_t & token, std::string & error); + +class Qwen3VLBridge { + public: + ~Qwen3VLBridge(); + + Qwen3VLBridge(const Qwen3VLBridge &) = delete; + Qwen3VLBridge & operator=(const Qwen3VLBridge &) = delete; + + static std::unique_ptr load(const Qwen3VLBridgeConfig & config, + std::string & error); + + bool extract_token_embeddings(const std::vector & images, + const std::string & instruction, int32_t token_id, + size_t token_count, std::vector & embeddings, + std::string & error); + + // Full conditioning sequence. Qwen3 uses the outer recorder's raw final + // decoder output (`l_out-(N-1)`); Qwen2.5 uses `result_norm`, matching its + // Transformers hidden_states[-1]. Values are widened from BF16. + bool extract_full_hidden_states(const std::vector & images, + const std::string & instruction, + std::vector & hidden_states, + std::vector & attention_mask, + std::string & error); + + // hidden_tuple_indices use the pinned Transformers 4.57 convention. Index + // zero is the embedding input and is not exposed. For Qwen3, in-place + // DeepStack aliases make the first D entries `deepstack_out`; remaining + // entries, including N, are raw `l_out`. For Qwen2.5, indices 1..N-1 are + // raw `l_out` and index N is `result_norm`. The result is layer-major + // [requested states, tokens, hidden size]. + bool extract_layer_hidden_states(const std::vector & images, + const std::string & instruction, + const std::vector & hidden_tuple_indices, + std::vector & hidden_states, + std::vector & attention_mask, + std::string & error); + + // Runs a full multimodal prefill followed by incremental KV-cached text + // decoding. The returned sequence includes the prompt, matching + // Transformers generate(return_dict_in_generate=false). + bool generate_autoregressive( + const std::vector & images, + const std::string & instruction, + const QwenVLGenerationConfig & generation, + QwenVLGenerationResult & result, std::string & error); + + void reset(); + const char * backend_name() const; + const char * text_attention_mode_name() const; + QwenVLArchitecture architecture() const; + + private: + struct Impl; + + explicit Qwen3VLBridge(std::unique_ptr impl); + + std::unique_ptr impl_; +}; + +// Neutral aliases for new callers. The original names remain the ABI/source +// compatibility surface for the completed Qwen3 integrations. +using QwenVLImageView = Qwen3VLImageView; +using QwenVLBridgeConfig = Qwen3VLBridgeConfig; +using QwenVLBridge = Qwen3VLBridge; + +} // namespace robotcpp::starvla diff --git a/tools/hf2gguf/README.md b/tools/hf2gguf/README.md index 57575c2..3cf2fc0 100644 --- a/tools/hf2gguf/README.md +++ b/tools/hf2gguf/README.md @@ -8,6 +8,7 @@ This directory contains tools for converting checkpoints to GGUF. - `smolvla/`: converts LeRobot-style SmolVLA checkpoints into four GGUF components. - `pi0/`: converts LeRobot-style pi0 checkpoints into six split GGUF components. +- `starvla/`: pinned Qwen3-VL and Qwen2.5-VL conversion plus a shared 3% CUDA action parity gate for seven StarVLA variants. No official finetuned Qwen3 FAST policy checkpoint is available. - `environment.yaml`: conda environment for the converters. ## Usage diff --git a/tools/hf2gguf/README_ZH.md b/tools/hf2gguf/README_ZH.md index 4219bff..f0fb6e4 100644 --- a/tools/hf2gguf/README_ZH.md +++ b/tools/hf2gguf/README_ZH.md @@ -8,6 +8,7 @@ - `smolvla/`：将 SmolVLA的lerobot-style的checkpoint 转成四个 GGUF component。 - `pi0/`：将 pi0的lerobot-style的checkpoint 转成六个 split GGUF component。 +- `starvla/`：固定版本的 Qwen3-VL/Qwen2.5-VL 转换与共享的 3% CUDA action parity gate，覆盖七种 StarVLA variant；上游尚无官方 finetuned Qwen3 FAST policy checkpoint。 - `environment.yaml`：converter conda 环境。 ## 使用说明 diff --git a/tools/hf2gguf/environment.yaml b/tools/hf2gguf/environment.yaml index a27464c..f014d70 100644 --- a/tools/hf2gguf/environment.yaml +++ b/tools/hf2gguf/environment.yaml @@ -8,6 +8,8 @@ dependencies: - torch - numpy - safetensors + - huggingface_hub>=0.36.0 - sentencepiece - transformers==4.53.2 + - pillow==12.1.1 - pyyaml diff --git a/tools/hf2gguf/starvla/__init__.py b/tools/hf2gguf/starvla/__init__.py new file mode 100755 index 0000000..52b10f1 --- /dev/null +++ b/tools/hf2gguf/starvla/__init__.py @@ -0,0 +1 @@ +"""StarVLA checkpoint conversion tools.""" diff --git a/tools/hf2gguf/starvla/compare_starvla_actions.py b/tools/hf2gguf/starvla/compare_starvla_actions.py new file mode 100644 index 0000000..e2422dc --- /dev/null +++ b/tools/hf2gguf/starvla/compare_starvla_actions.py @@ -0,0 +1,82 @@ +#!/usr/bin/env python3 +"""Compare C++ StarVLA actions with a local Python reference.""" + +from __future__ import annotations + +import argparse +import json +import math +import sys +from pathlib import Path +from typing import Any + +import numpy as np + + +ACTION_RELATIVE_L2_LIMIT = 0.03 + + +class ComparisonError(RuntimeError): + pass + + +def load_actions(path: Path, key: str) -> np.ndarray: + try: + value: Any = json.loads(path.read_text(encoding="utf-8")) + for part in key.split("."): + value = value[part] + actions = np.asarray(value, dtype=np.float64) + except (OSError, json.JSONDecodeError, KeyError, TypeError, ValueError) as exc: + raise ComparisonError(f"cannot load {key!r} from {path}: {exc}") from exc + if actions.ndim == 3 and actions.shape[0] == 1: + actions = actions[0] + if actions.ndim != 2 or 0 in actions.shape: + raise ComparisonError(f"{path}:{key} must have shape [steps, dims]") + if not np.isfinite(actions).all(): + raise ComparisonError(f"{path}:{key} contains non-finite values") + return actions + + +def compare_actions(reference: np.ndarray, candidate: np.ndarray) -> dict[str, Any]: + if reference.shape != candidate.shape: + raise ComparisonError( + f"action shape mismatch: reference={reference.shape}, candidate={candidate.shape}" + ) + difference_l2 = float(np.linalg.norm(candidate - reference)) + reference_l2 = float(np.linalg.norm(reference)) + relative_l2 = difference_l2 / reference_l2 if reference_l2 else ( + 0.0 if difference_l2 == 0.0 else math.inf + ) + return { + "shape": list(reference.shape), + "relative_l2": relative_l2, + "limit": ACTION_RELATIVE_L2_LIMIT, + "passed": relative_l2 <= ACTION_RELATIVE_L2_LIMIT + 1e-12, + } + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--reference", type=Path, required=True) + parser.add_argument("--candidate", type=Path, required=True) + parser.add_argument("--reference-key", default="outputs.unnormalized_actions") + parser.add_argument("--candidate-key", default="actions") + return parser.parse_args() + + +def main() -> int: + args = parse_args() + try: + result = compare_actions( + load_actions(args.reference, args.reference_key), + load_actions(args.candidate, args.candidate_key), + ) + except ComparisonError as exc: + print(f"error: {exc}", file=sys.stderr) + return 2 + print(json.dumps(result, indent=2)) + return 0 if result["passed"] else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/hf2gguf/starvla/convert_starvla_policy_to_gguf.py b/tools/hf2gguf/starvla/convert_starvla_policy_to_gguf.py new file mode 100755 index 0000000..c8c53ed --- /dev/null +++ b/tools/hf2gguf/starvla/convert_starvla_policy_to_gguf.py @@ -0,0 +1,2249 @@ +#!/usr/bin/env python3 +"""Convert a StarVLA policy staging directory to a policy GGUF.""" + +from __future__ import annotations + +import argparse +import copy +import hashlib +import json +import math +import os +import sys +from pathlib import Path +from typing import Any + +import numpy as np + +from starvla_checkpoint import ( + DEFAULT_CATALOG, + DEFAULT_MMPROJ_DTYPE, + DEFAULT_POLICY_DTYPE, + DEFAULT_TEXT_DTYPE, + StarVLAError, + create_output_temporary, + default_mmproj_filename, + default_text_filename, + get_variant, + load_catalog, + resolve_effective_config, + sha256_file, + validate_official_surgery_manifest, + verify_staged_assets, + verify_staged_tensors_against_checkpoint, +) + + +OFT_TENSOR_MAP = { + "action_model.model.layer_norm1.weight": "starvla.policy.oft.input_norm.weight", + "action_model.model.layer_norm1.bias": "starvla.policy.oft.input_norm.bias", + "action_model.model.fc1.weight": "starvla.policy.oft.input_proj.weight", + "action_model.model.fc1.bias": "starvla.policy.oft.input_proj.bias", + "action_model.model.mlp_resnet_blocks.0.ffn.0.weight": "starvla.policy.oft.block.0.norm.weight", + "action_model.model.mlp_resnet_blocks.0.ffn.0.bias": "starvla.policy.oft.block.0.norm.bias", + "action_model.model.mlp_resnet_blocks.0.ffn.1.weight": "starvla.policy.oft.block.0.linear.weight", + "action_model.model.mlp_resnet_blocks.0.ffn.1.bias": "starvla.policy.oft.block.0.linear.bias", + "action_model.model.mlp_resnet_blocks.1.ffn.0.weight": "starvla.policy.oft.block.1.norm.weight", + "action_model.model.mlp_resnet_blocks.1.ffn.0.bias": "starvla.policy.oft.block.1.norm.bias", + "action_model.model.mlp_resnet_blocks.1.ffn.1.weight": "starvla.policy.oft.block.1.linear.weight", + "action_model.model.mlp_resnet_blocks.1.ffn.1.bias": "starvla.policy.oft.block.1.linear.bias", + "action_model.model.layer_norm2.weight": "starvla.policy.oft.output_norm.weight", + "action_model.model.layer_norm2.bias": "starvla.policy.oft.output_norm.bias", + "action_model.model.fc2.weight": "starvla.policy.oft.output_proj.weight", + "action_model.model.fc2.bias": "starvla.policy.oft.output_proj.bias", +} + + +def build_groot_tensor_map(block_count: int = 16) -> dict[str, str]: + """Return the complete released Qwen-GR00T policy tensor renaming map.""" + tensor_map = { + "action_model.model.timestep_encoder.timestep_embedder.linear_1.weight": + "starvla.policy.groot.timestep.input.weight", + "action_model.model.timestep_encoder.timestep_embedder.linear_1.bias": + "starvla.policy.groot.timestep.input.bias", + "action_model.model.timestep_encoder.timestep_embedder.linear_2.weight": + "starvla.policy.groot.timestep.output.weight", + "action_model.model.timestep_encoder.timestep_embedder.linear_2.bias": + "starvla.policy.groot.timestep.output.bias", + } + block_suffixes = { + "norm1.linear.weight": "ada_norm.weight", + "norm1.linear.bias": "ada_norm.bias", + "attn1.to_q.weight": "attention.query.weight", + "attn1.to_q.bias": "attention.query.bias", + "attn1.to_k.weight": "attention.key.weight", + "attn1.to_k.bias": "attention.key.bias", + "attn1.to_v.weight": "attention.value.weight", + "attn1.to_v.bias": "attention.value.bias", + "attn1.to_out.0.weight": "attention.output.weight", + "attn1.to_out.0.bias": "attention.output.bias", + "ff.net.0.proj.weight": "feed_forward.input.weight", + "ff.net.0.proj.bias": "feed_forward.input.bias", + "ff.net.2.weight": "feed_forward.output.weight", + "ff.net.2.bias": "feed_forward.output.bias", + } + for block in range(block_count): + for source_suffix, destination_suffix in block_suffixes.items(): + tensor_map[f"action_model.model.transformer_blocks.{block}.{source_suffix}"] = ( + f"starvla.policy.groot.block.{block}.{destination_suffix}" + ) + tensor_map.update( + { + "action_model.model.proj_out_1.weight": "starvla.policy.groot.output.modulation.weight", + "action_model.model.proj_out_1.bias": "starvla.policy.groot.output.modulation.bias", + "action_model.model.proj_out_2.weight": "starvla.policy.groot.output.projection.weight", + "action_model.model.proj_out_2.bias": "starvla.policy.groot.output.projection.bias", + "action_model.action_encoder.layer1.weight": "starvla.policy.groot.action.input.weight", + "action_model.action_encoder.layer1.bias": "starvla.policy.groot.action.input.bias", + "action_model.action_encoder.layer2.weight": "starvla.policy.groot.action.time_mix.weight", + "action_model.action_encoder.layer2.bias": "starvla.policy.groot.action.time_mix.bias", + "action_model.action_encoder.layer3.weight": "starvla.policy.groot.action.output.weight", + "action_model.action_encoder.layer3.bias": "starvla.policy.groot.action.output.bias", + "action_model.action_decoder.layer1.weight": "starvla.policy.groot.velocity.input.weight", + "action_model.action_decoder.layer1.bias": "starvla.policy.groot.velocity.input.bias", + "action_model.action_decoder.layer2.weight": "starvla.policy.groot.velocity.output.weight", + "action_model.action_decoder.layer2.bias": "starvla.policy.groot.velocity.output.bias", + "action_model.future_tokens.weight": "starvla.policy.groot.future_tokens.weight", + "action_model.position_embedding.weight": "starvla.policy.groot.action_position.weight", + } + ) + return tensor_map + + +def build_pi_tensor_map(block_count: int = 16) -> dict[str, str]: + """Return tensors used by the legacy Qwen-PI inference graph.""" + tensor_map = { + "action_model.model.timestep_encoder.timestep_embedder.linear_1.weight": + "starvla.policy.pi.timestep.input.weight", + "action_model.model.timestep_encoder.timestep_embedder.linear_1.bias": + "starvla.policy.pi.timestep.input.bias", + "action_model.model.timestep_encoder.timestep_embedder.linear_2.weight": + "starvla.policy.pi.timestep.output.weight", + "action_model.model.timestep_encoder.timestep_embedder.linear_2.bias": + "starvla.policy.pi.timestep.output.bias", + } + block_suffixes = { + "norm1.linear.weight": "ada_norm.weight", + "norm1.linear.bias": "ada_norm.bias", + "attn1.to_q.weight": "attention.query.weight", + "attn1.to_q.bias": "attention.query.bias", + "attn1.to_k.weight": "attention.key.weight", + "attn1.to_k.bias": "attention.key.bias", + "attn1.to_v.weight": "attention.value.weight", + "attn1.to_v.bias": "attention.value.bias", + "attn1.to_out.0.weight": "attention.output.weight", + "attn1.to_out.0.bias": "attention.output.bias", + "ff.net.0.proj.weight": "feed_forward.input.weight", + "ff.net.0.proj.bias": "feed_forward.input.bias", + "ff.net.2.weight": "feed_forward.output.weight", + "ff.net.2.bias": "feed_forward.output.bias", + } + for block in range(block_count): + for source_suffix, destination_suffix in block_suffixes.items(): + tensor_map[f"action_model.model.transformer_blocks.{block}.{source_suffix}"] = ( + f"starvla.policy.pi.block.{block}.{destination_suffix}" + ) + tensor_map.update( + { + "action_model.state_encoder.layer1.weight": "starvla.policy.pi.state.input.weight", + "action_model.state_encoder.layer1.bias": "starvla.policy.pi.state.input.bias", + "action_model.state_encoder.layer2.weight": "starvla.policy.pi.state.output.weight", + "action_model.state_encoder.layer2.bias": "starvla.policy.pi.state.output.bias", + "action_model.action_encoder.layer1.weight": "starvla.policy.pi.action.input.weight", + "action_model.action_encoder.layer1.bias": "starvla.policy.pi.action.input.bias", + "action_model.action_encoder.layer2.weight": + "starvla.policy.pi.action.time_mix.weight", + "action_model.action_encoder.layer2.bias": + "starvla.policy.pi.action.time_mix.bias", + "action_model.action_encoder.layer3.weight": "starvla.policy.pi.action.output.weight", + "action_model.action_encoder.layer3.bias": "starvla.policy.pi.action.output.bias", + "action_model.action_decoder.layer1.weight": + "starvla.policy.pi.velocity.input.weight", + "action_model.action_decoder.layer1.bias": "starvla.policy.pi.velocity.input.bias", + "action_model.action_decoder.layer2.weight": + "starvla.policy.pi.velocity.output.weight", + "action_model.action_decoder.layer2.bias": "starvla.policy.pi.velocity.output.bias", + "action_model.future_tokens.weight": "starvla.policy.pi.future_tokens.weight", + "action_model.position_embedding.weight": "starvla.policy.pi.action_position.weight", + } + ) + return tensor_map + + +def build_pi_v3_tensor_map( + block_count: int = 36, + projector_count: int = 36, +) -> dict[str, str]: + """Return the tensors used by the Qwen PI-v3 inference graph.""" + tensor_map = { + "action_model.model.timestep_encoder.timestep_embedder.linear_1.weight": + "starvla.policy.pi_v3.timestep.input.weight", + "action_model.model.timestep_encoder.timestep_embedder.linear_1.bias": + "starvla.policy.pi_v3.timestep.input.bias", + "action_model.model.timestep_encoder.timestep_embedder.linear_2.weight": + "starvla.policy.pi_v3.timestep.output.weight", + "action_model.model.timestep_encoder.timestep_embedder.linear_2.bias": + "starvla.policy.pi_v3.timestep.output.bias", + } + block_suffixes = { + "norm1.linear.weight": "ada_norm.weight", + "norm1.linear.bias": "ada_norm.bias", + "attn1.to_q.weight": "attention.query.weight", + "attn1.to_q.bias": "attention.query.bias", + "attn1.to_k.weight": "attention.key.weight", + "attn1.to_k.bias": "attention.key.bias", + "attn1.to_v.weight": "attention.value.weight", + "attn1.to_v.bias": "attention.value.bias", + "attn1.to_out.0.weight": "attention.output.weight", + "attn1.to_out.0.bias": "attention.output.bias", + "ff.net.0.proj.weight": "feed_forward.input.weight", + "ff.net.0.proj.bias": "feed_forward.input.bias", + "ff.net.2.weight": "feed_forward.output.weight", + "ff.net.2.bias": "feed_forward.output.bias", + } + for block in range(block_count): + for source_suffix, destination_suffix in block_suffixes.items(): + tensor_map[f"action_model.model.transformer_blocks.{block}.{source_suffix}"] = ( + f"starvla.policy.pi_v3.block.{block}.{destination_suffix}" + ) + tensor_map.update( + { + "action_model.action_encoder.layer1.weight": "starvla.policy.pi_v3.action.input.weight", + "action_model.action_encoder.layer1.bias": "starvla.policy.pi_v3.action.input.bias", + "action_model.action_encoder.layer2.weight": "starvla.policy.pi_v3.action.time_mix.weight", + "action_model.action_encoder.layer2.bias": "starvla.policy.pi_v3.action.time_mix.bias", + "action_model.action_encoder.layer3.weight": "starvla.policy.pi_v3.action.output.weight", + "action_model.action_encoder.layer3.bias": "starvla.policy.pi_v3.action.output.bias", + "action_model.action_decoder.layer1.weight": "starvla.policy.pi_v3.velocity.input.weight", + "action_model.action_decoder.layer1.bias": "starvla.policy.pi_v3.velocity.input.bias", + "action_model.action_decoder.layer2.weight": "starvla.policy.pi_v3.velocity.output.weight", + "action_model.action_decoder.layer2.bias": "starvla.policy.pi_v3.velocity.output.bias", + "action_model.future_tokens.weight": "starvla.policy.pi_v3.future_tokens.weight", + "action_model.position_embedding.weight": "starvla.policy.pi_v3.action_position.weight", + } + ) + for projector in range(projector_count): + source_prefix = f"project_layers.{projector}" + destination_prefix = f"starvla.policy.pi_v3.projector.{projector}" + tensor_map.update( + { + f"{source_prefix}.0.weight": f"{destination_prefix}.norm.weight", + f"{source_prefix}.0.bias": f"{destination_prefix}.norm.bias", + f"{source_prefix}.1.weight": f"{destination_prefix}.projection.weight", + f"{source_prefix}.1.bias": f"{destination_prefix}.projection.bias", + } + ) + return tensor_map + + +GROOT_BLOCK_COUNT = 16 +GROOT_TENSOR_MAP = build_groot_tensor_map(GROOT_BLOCK_COUNT) +GROOT_UNUSED_SOURCE_TENSORS = { + "action_model.state_encoder.layer1.weight", + "action_model.state_encoder.layer1.bias", + "action_model.state_encoder.layer2.weight", + "action_model.state_encoder.layer2.bias", +} +GROOT_SOURCE_TENSOR_COUNT = 248 +GROOT_POLICY_TENSOR_COUNT = 244 +GROOT_QWEN3_POLICY_NUMEL = 161_472_775 +GROOT_QWEN25_POLICY_NUMEL = 155_181_319 +GROOT_DIT_NORM_EPS = 1e-5 +GROOT_OUTPUT_NORM_EPS = 1e-6 +GROOT_OFFICIAL_DIMENSIONS_BY_BACKBONE = { + backbone: { + "qwen_hidden_dim": qwen_hidden_dim, + "dit_width": 768, + "timestep_dim": 256, + "feed_forward_dim": 3072, + "output_dim": 1024, + "mlp_hidden_dim": 1024, + "state_dim": 7, + "action_dim": 7, + "future_token_count": 32, + "max_sequence_length": 1024, + "block_count": GROOT_BLOCK_COUNT, + "tensor_count": GROOT_SOURCE_TENSOR_COUNT, + "numel": numel, + } + for backbone, qwen_hidden_dim, numel in ( + ("qwen3_vl", 2560, GROOT_QWEN3_POLICY_NUMEL), + ("qwen2_5_vl", 2048, GROOT_QWEN25_POLICY_NUMEL), + ) +} +GROOT_OFFICIAL_DIMENSIONS = GROOT_OFFICIAL_DIMENSIONS_BY_BACKBONE["qwen3_vl"] + +PI_BLOCK_COUNT = 16 +PI_TENSOR_MAP = build_pi_tensor_map(PI_BLOCK_COUNT) +PI_UNUSED_SOURCE_TENSORS = { + "action_model.model.proj_out_1.weight", + "action_model.model.proj_out_1.bias", + "action_model.model.proj_out_2.weight", + "action_model.model.proj_out_2.bias", +} +PI_POLICY_TENSOR_COUNT = 244 +PI_POLICY_NUMEL = 967_796_743 +PI_DIT_NORM_EPS = 1e-5 +PI_OFFICIAL_DIMENSIONS = { + "qwen_hidden_dim": 2048, + "dit_width": 2048, + "timestep_dim": 256, + "feed_forward_dim": 8192, + "mlp_hidden_dim": 2048, + "state_dim": 7, + "action_dim": 7, + "future_token_count": 32, + "max_sequence_length": 1024, + "block_count": PI_BLOCK_COUNT, + "tensor_count": PI_POLICY_TENSOR_COUNT, + "numel": PI_POLICY_NUMEL, +} + +PI_V3_BLOCK_COUNT = 36 +PI_V3_PROJECTOR_COUNT = 36 +PI_V3_TENSOR_MAP = build_pi_v3_tensor_map(PI_V3_BLOCK_COUNT, PI_V3_PROJECTOR_COUNT) +PI_V3_POLICY_TENSOR_COUNT = len(PI_V3_TENSOR_MAP) +PI_V3_DIT_NORM_EPS = 1e-5 +PI_V3_PROJECTOR_NORM_EPS = 1e-5 +PI_V3_OFFICIAL_DIMENSIONS = { + "qwen_hidden_dim": 2560, + "dit_width": 1024, + "timestep_dim": 256, + "feed_forward_dim": 4096, + "mlp_hidden_dim": 1024, + "action_dim": 7, + "future_token_count": 32, + "max_sequence_length": 1024, + "block_count": PI_V3_BLOCK_COUNT, + "projector_count": PI_V3_PROJECTOR_COUNT, + "tensor_count": PI_V3_POLICY_TENSOR_COUNT, +} + +ACTION_NAMES = ["x", "y", "z", "roll", "pitch", "yaw", "gripper"] +OFT_ACTION_TOKEN = "🔍" +OFT_ACTION_TOKEN_ID = 146663 +OFT_LAYER_NORM_EPS = 1e-5 +QWEN3VL_PROCESSOR_MIN_PIXELS = 65_536 +QWEN3VL_PROCESSOR_MAX_PIXELS = 16_777_216 +QWEN3VL_IMAGE_PATCH_SIZE = 16 +QWEN3VL_TEMPORAL_PATCH_SIZE = 2 +QWEN3VL_SPATIAL_MERGE_SIZE = 2 +QWEN3VL_MIN_IMAGE_TOKENS = 64 +QWEN3VL_MAX_IMAGE_TOKENS = 16_384 +QWEN3VL_IMAGE_MEAN = [0.5, 0.5, 0.5] +QWEN3VL_IMAGE_STD = [0.5, 0.5, 0.5] +QWEN25VL_PROCESSOR_MIN_PIXELS = 3_136 +QWEN25VL_PROCESSOR_MAX_PIXELS = 12_845_056 +QWEN25VL_IMAGE_PATCH_SIZE = 14 +QWEN25VL_TEMPORAL_PATCH_SIZE = 2 +QWEN25VL_SPATIAL_MERGE_SIZE = 2 +QWEN25VL_MIN_IMAGE_TOKENS = 4 +QWEN25VL_MAX_IMAGE_TOKENS = 16_384 +QWEN25VL_IMAGE_MEAN = [0.48145466, 0.4578275, 0.40821073] +QWEN25VL_IMAGE_STD = [0.26862954, 0.26130258, 0.27577711] +# These defaults are executable behavior in the pinned Transformers 4.57 fast processor, +# including antialias=True on its torchvision resize call. +QWEN3VL_DYNAMIC_IMAGE_METADATA = { + "starvla.image.count": 1, + "starvla.image.names": ["image_0"], + "starvla.image.preprocessing_mode": "qwen3vl_smart_resize", + "starvla.image.framework_inference_pre_resize": False, + "starvla.image.framework_inference_pre_resize_config_key": ( + "datasets.vla_data.obs_image_size" + ), + "starvla.image.processor_min_pixels": QWEN3VL_PROCESSOR_MIN_PIXELS, + "starvla.image.processor_max_pixels": QWEN3VL_PROCESSOR_MAX_PIXELS, + "starvla.image.processor_class": "Qwen2VLImageProcessorFast", + "starvla.image.processor_reference_transformers_version": "4.57.0", + "starvla.image.processor_do_convert_rgb": True, + "starvla.image.processor_do_resize": True, + "starvla.image.processor_resize_resample": "bicubic", + "starvla.image.processor_resize_antialias": True, + "starvla.image.processor_do_rescale": True, + "starvla.image.processor_rescale_factor": 1.0 / 255.0, + "starvla.image.processor_do_normalize": True, + "starvla.image.processor_image_mean": QWEN3VL_IMAGE_MEAN, + "starvla.image.processor_image_std": QWEN3VL_IMAGE_STD, + "starvla.image.patch_size": QWEN3VL_IMAGE_PATCH_SIZE, + "starvla.image.temporal_patch_size": QWEN3VL_TEMPORAL_PATCH_SIZE, + "starvla.image.spatial_merge_size": QWEN3VL_SPATIAL_MERGE_SIZE, + "starvla.image.token_count_mode": "dynamic_grid_thw_after_spatial_merge", + "starvla.image.min_token_count": QWEN3VL_MIN_IMAGE_TOKENS, + "starvla.image.max_token_count": QWEN3VL_MAX_IMAGE_TOKENS, +} + + +def _load_json(path: Path) -> dict[str, Any]: + try: + value = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise StarVLAError(f"failed to load JSON {path}: {exc}") from exc + if not isinstance(value, dict): + raise StarVLAError(f"expected a JSON object in {path}") + return value + + +def _write_gguf_arrays_no_overwrite( + output: Path, + metadata: dict[str, Any], + arrays: Any, + writer: Any, +) -> None: + """Write beside the destination, then publish without replacing an existing file.""" + output.parent.mkdir(parents=True, exist_ok=True) + if output.exists() or output.is_symlink(): + raise StarVLAError(f"refusing to overwrite existing output: {output}") + + descriptor, temporary = create_output_temporary(output) + os.close(descriptor) + try: + writer(temporary, metadata, arrays) + if not temporary.is_file() or temporary.stat().st_size == 0: + raise StarVLAError(f"GGUF writer did not create a non-empty output: {temporary}") + try: + os.link(temporary, output) + except FileExistsError as exc: + raise StarVLAError(f"refusing to overwrite existing output: {output}") from exc + temporary.unlink() + finally: + temporary.unlink(missing_ok=True) + + +def _load_yaml(path: Path) -> dict[str, Any]: + try: + import yaml + except ImportError as exc: + raise StarVLAError("PyYAML is required to load a StarVLA policy config") from exc + try: + value = yaml.safe_load(path.read_text(encoding="utf-8")) + except (OSError, yaml.YAMLError) as exc: + raise StarVLAError(f"failed to load YAML {path}: {exc}") from exc + if not isinstance(value, dict): + raise StarVLAError(f"expected a YAML object in {path}") + return value + + +def load_policy_tensors(policy_dir: Path) -> dict[str, Any]: + try: + from safetensors import safe_open + except ImportError as exc: + raise StarVLAError("safetensors is required to convert a StarVLA policy") from exc + + index_path = policy_dir / "policy.safetensors.index.json" + index = _load_json(index_path) + weight_map = index.get("weight_map") + if not isinstance(weight_map, dict) or not weight_map: + raise StarVLAError(f"invalid or empty safetensors weight_map in {index_path}") + + tensors = {} + by_shard: dict[str, list[str]] = {} + for name, shard in weight_map.items(): + by_shard.setdefault(str(shard), []).append(str(name)) + for shard, names in sorted(by_shard.items()): + shard_path = policy_dir / shard + if not shard_path.is_file(): + raise StarVLAError(f"missing policy shard: {shard_path}") + with safe_open(shard_path, framework="pt", device="cpu") as handle: + if set(handle.keys()) != set(names): + raise StarVLAError(f"policy shard/index key mismatch: {shard_path}") + for name in sorted(names): + tensors[name] = handle.get_tensor(name) + if set(tensors) != set(weight_map): + raise StarVLAError("loaded policy tensor set does not match the index") + return tensors + + +def validate_oft_tensors(tensors: dict[str, Any]) -> dict[str, int]: + actual = set(tensors) + expected = set(OFT_TENSOR_MAP) + if actual != expected: + missing = sorted(expected - actual) + unexpected = sorted(actual - expected) + raise StarVLAError(f"OFT policy tensor mismatch; missing={missing}, unexpected={unexpected}") + + def shape(name: str) -> list[int]: + return [int(dim) for dim in tensors[name].shape] + + input_dim = shape("action_model.model.layer_norm1.weight")[0] + input_projection = shape("action_model.model.fc1.weight") + if len(input_projection) != 2 or input_projection[1] != input_dim: + raise StarVLAError(f"invalid OFT input projection shape: {input_projection}") + hidden_dim = input_projection[0] + output_projection = shape("action_model.model.fc2.weight") + if len(output_projection) != 2 or output_projection[1] != hidden_dim: + raise StarVLAError(f"invalid OFT output projection shape: {output_projection}") + action_dim = output_projection[0] + + expected_shapes = { + "action_model.model.layer_norm1.weight": [input_dim], + "action_model.model.layer_norm1.bias": [input_dim], + "action_model.model.fc1.weight": [hidden_dim, input_dim], + "action_model.model.fc1.bias": [hidden_dim], + "action_model.model.mlp_resnet_blocks.0.ffn.0.weight": [hidden_dim], + "action_model.model.mlp_resnet_blocks.0.ffn.0.bias": [hidden_dim], + "action_model.model.mlp_resnet_blocks.0.ffn.1.weight": [hidden_dim, hidden_dim], + "action_model.model.mlp_resnet_blocks.0.ffn.1.bias": [hidden_dim], + "action_model.model.mlp_resnet_blocks.1.ffn.0.weight": [hidden_dim], + "action_model.model.mlp_resnet_blocks.1.ffn.0.bias": [hidden_dim], + "action_model.model.mlp_resnet_blocks.1.ffn.1.weight": [hidden_dim, hidden_dim], + "action_model.model.mlp_resnet_blocks.1.ffn.1.bias": [hidden_dim], + "action_model.model.layer_norm2.weight": [hidden_dim], + "action_model.model.layer_norm2.bias": [hidden_dim], + "action_model.model.fc2.weight": [action_dim, hidden_dim], + "action_model.model.fc2.bias": [action_dim], + } + mismatches = [ + f"{name}: expected {expected_shape}, got {shape(name)}" + for name, expected_shape in expected_shapes.items() + if shape(name) != expected_shape + ] + if mismatches: + raise StarVLAError("invalid OFT tensor shapes: " + "; ".join(mismatches)) + return {"input_dim": input_dim, "hidden_dim": hidden_dim, "action_dim": action_dim} + + +def validate_groot_tensors(tensors: dict[str, Any]) -> dict[str, int]: + """Validate every released GR00T policy tensor and infer its architecture.""" + actual = set(tensors) + expected = set(GROOT_TENSOR_MAP) + if not expected.issubset(actual) or actual - expected != GROOT_UNUSED_SOURCE_TENSORS: + missing = sorted(expected - actual) + unexpected = sorted(actual - expected - GROOT_UNUSED_SOURCE_TENSORS) + raise StarVLAError(f"GR00T policy tensor mismatch; missing={missing}, unexpected={unexpected}") + + def shape(name: str) -> list[int]: + return [int(dim) for dim in tensors[name].shape] + + def matrix_shape(name: str) -> list[int]: + value = shape(name) + if len(value) != 2: + raise StarVLAError(f"invalid GR00T matrix shape for {name}: {value}") + return value + + timestep_input = matrix_shape( + "action_model.model.timestep_encoder.timestep_embedder.linear_1.weight" + ) + dit_width, timestep_dim = timestep_input + cross_attention_dim = matrix_shape( + "action_model.model.transformer_blocks.0.attn1.to_k.weight" + )[1] + feed_forward_dim = matrix_shape( + "action_model.model.transformer_blocks.0.ff.net.0.proj.weight" + )[0] + output_dim = matrix_shape("action_model.model.proj_out_2.weight")[0] + mlp_hidden_dim = matrix_shape("action_model.action_decoder.layer1.weight")[0] + state_dim = matrix_shape("action_model.state_encoder.layer1.weight")[1] + action_dim = matrix_shape("action_model.action_encoder.layer1.weight")[1] + future_token_count = matrix_shape("action_model.future_tokens.weight")[0] + max_sequence_length = matrix_shape("action_model.position_embedding.weight")[0] + + expected_shapes = { + "action_model.model.timestep_encoder.timestep_embedder.linear_1.weight": [ + dit_width, + timestep_dim, + ], + "action_model.model.timestep_encoder.timestep_embedder.linear_1.bias": [dit_width], + "action_model.model.timestep_encoder.timestep_embedder.linear_2.weight": [ + dit_width, + dit_width, + ], + "action_model.model.timestep_encoder.timestep_embedder.linear_2.bias": [dit_width], + "action_model.model.proj_out_1.weight": [2 * dit_width, dit_width], + "action_model.model.proj_out_1.bias": [2 * dit_width], + "action_model.model.proj_out_2.weight": [output_dim, dit_width], + "action_model.model.proj_out_2.bias": [output_dim], + "action_model.action_encoder.layer1.weight": [dit_width, action_dim], + "action_model.action_encoder.layer1.bias": [dit_width], + "action_model.action_encoder.layer2.weight": [dit_width, 2 * dit_width], + "action_model.action_encoder.layer2.bias": [dit_width], + "action_model.action_encoder.layer3.weight": [dit_width, dit_width], + "action_model.action_encoder.layer3.bias": [dit_width], + "action_model.action_decoder.layer1.weight": [mlp_hidden_dim, output_dim], + "action_model.action_decoder.layer1.bias": [mlp_hidden_dim], + "action_model.action_decoder.layer2.weight": [action_dim, mlp_hidden_dim], + "action_model.action_decoder.layer2.bias": [action_dim], + "action_model.future_tokens.weight": [future_token_count, dit_width], + "action_model.position_embedding.weight": [max_sequence_length, dit_width], + } + for block in range(GROOT_BLOCK_COUNT): + prefix = f"action_model.model.transformer_blocks.{block}" + attention_input_dim = cross_attention_dim if block % 2 == 0 else dit_width + expected_shapes.update( + { + f"{prefix}.norm1.linear.weight": [2 * dit_width, dit_width], + f"{prefix}.norm1.linear.bias": [2 * dit_width], + f"{prefix}.attn1.to_q.weight": [dit_width, dit_width], + f"{prefix}.attn1.to_q.bias": [dit_width], + f"{prefix}.attn1.to_k.weight": [dit_width, attention_input_dim], + f"{prefix}.attn1.to_k.bias": [dit_width], + f"{prefix}.attn1.to_v.weight": [dit_width, attention_input_dim], + f"{prefix}.attn1.to_v.bias": [dit_width], + f"{prefix}.attn1.to_out.0.weight": [dit_width, dit_width], + f"{prefix}.attn1.to_out.0.bias": [dit_width], + f"{prefix}.ff.net.0.proj.weight": [feed_forward_dim, dit_width], + f"{prefix}.ff.net.0.proj.bias": [feed_forward_dim], + f"{prefix}.ff.net.2.weight": [dit_width, feed_forward_dim], + f"{prefix}.ff.net.2.bias": [dit_width], + } + ) + mismatches = [ + f"{name}: expected {expected_shape}, got {shape(name)}" + for name, expected_shape in expected_shapes.items() + if shape(name) != expected_shape + ] + if mismatches: + raise StarVLAError("invalid GR00T tensor shapes: " + "; ".join(mismatches)) + + numel = sum(int(tensor.numel()) for tensor in tensors.values()) + return { + "qwen_hidden_dim": cross_attention_dim, + "dit_width": dit_width, + "timestep_dim": timestep_dim, + "feed_forward_dim": feed_forward_dim, + "output_dim": output_dim, + "mlp_hidden_dim": mlp_hidden_dim, + "state_dim": state_dim, + "action_dim": action_dim, + "future_token_count": future_token_count, + "max_sequence_length": max_sequence_length, + "block_count": GROOT_BLOCK_COUNT, + "tensor_count": len(tensors), + "numel": numel, + } + + +def validate_pi_tensors(tensors: dict[str, Any]) -> dict[str, int]: + """Validate the tensors used by the legacy Qwen-PI inference graph.""" + actual = set(tensors) + expected = set(PI_TENSOR_MAP) + if not expected.issubset(actual) or actual - expected != PI_UNUSED_SOURCE_TENSORS: + missing = sorted(expected - actual) + unexpected = sorted(actual - expected - PI_UNUSED_SOURCE_TENSORS) + raise StarVLAError( + f"legacy PI policy tensor mismatch; missing={missing}, unexpected={unexpected}" + ) + + def shape(name: str) -> list[int]: + return [int(dim) for dim in tensors[name].shape] + + def matrix_shape(name: str) -> list[int]: + value = shape(name) + if len(value) != 2: + raise StarVLAError(f"invalid legacy PI matrix shape for {name}: {value}") + return value + + timestep_input = matrix_shape( + "action_model.model.timestep_encoder.timestep_embedder.linear_1.weight" + ) + dit_width, timestep_dim = timestep_input + cross_attention_dim = matrix_shape( + "action_model.model.transformer_blocks.0.attn1.to_k.weight" + )[1] + feed_forward_dim = matrix_shape( + "action_model.model.transformer_blocks.0.ff.net.0.proj.weight" + )[0] + mlp_hidden_dim, state_dim = matrix_shape("action_model.state_encoder.layer1.weight") + action_dim = matrix_shape("action_model.action_encoder.layer1.weight")[1] + future_token_count = matrix_shape("action_model.future_tokens.weight")[0] + max_sequence_length = matrix_shape("action_model.position_embedding.weight")[0] + + expected_shapes = { + "action_model.model.timestep_encoder.timestep_embedder.linear_1.weight": [ + dit_width, + timestep_dim, + ], + "action_model.model.timestep_encoder.timestep_embedder.linear_1.bias": [dit_width], + "action_model.model.timestep_encoder.timestep_embedder.linear_2.weight": [ + dit_width, + dit_width, + ], + "action_model.model.timestep_encoder.timestep_embedder.linear_2.bias": [dit_width], + "action_model.state_encoder.layer1.weight": [mlp_hidden_dim, state_dim], + "action_model.state_encoder.layer1.bias": [mlp_hidden_dim], + "action_model.state_encoder.layer2.weight": [dit_width, mlp_hidden_dim], + "action_model.state_encoder.layer2.bias": [dit_width], + "action_model.action_encoder.layer1.weight": [dit_width, action_dim], + "action_model.action_encoder.layer1.bias": [dit_width], + "action_model.action_encoder.layer2.weight": [dit_width, 2 * dit_width], + "action_model.action_encoder.layer2.bias": [dit_width], + "action_model.action_encoder.layer3.weight": [dit_width, dit_width], + "action_model.action_encoder.layer3.bias": [dit_width], + "action_model.action_decoder.layer1.weight": [mlp_hidden_dim, dit_width], + "action_model.action_decoder.layer1.bias": [mlp_hidden_dim], + "action_model.action_decoder.layer2.weight": [action_dim, mlp_hidden_dim], + "action_model.action_decoder.layer2.bias": [action_dim], + "action_model.future_tokens.weight": [future_token_count, dit_width], + "action_model.position_embedding.weight": [max_sequence_length, dit_width], + } + for block in range(PI_BLOCK_COUNT): + prefix = f"action_model.model.transformer_blocks.{block}" + expected_shapes.update( + { + f"{prefix}.norm1.linear.weight": [2 * dit_width, dit_width], + f"{prefix}.norm1.linear.bias": [2 * dit_width], + f"{prefix}.attn1.to_q.weight": [dit_width, dit_width], + f"{prefix}.attn1.to_q.bias": [dit_width], + f"{prefix}.attn1.to_k.weight": [dit_width, cross_attention_dim], + f"{prefix}.attn1.to_k.bias": [dit_width], + f"{prefix}.attn1.to_v.weight": [dit_width, cross_attention_dim], + f"{prefix}.attn1.to_v.bias": [dit_width], + f"{prefix}.attn1.to_out.0.weight": [dit_width, dit_width], + f"{prefix}.attn1.to_out.0.bias": [dit_width], + f"{prefix}.ff.net.0.proj.weight": [feed_forward_dim, dit_width], + f"{prefix}.ff.net.0.proj.bias": [feed_forward_dim], + f"{prefix}.ff.net.2.weight": [dit_width, feed_forward_dim], + f"{prefix}.ff.net.2.bias": [dit_width], + } + ) + mismatches = [ + f"{name}: expected {expected_shape}, got {shape(name)}" + for name, expected_shape in expected_shapes.items() + if shape(name) != expected_shape + ] + if mismatches: + raise StarVLAError("invalid legacy PI tensor shapes: " + "; ".join(mismatches)) + + return { + "qwen_hidden_dim": cross_attention_dim, + "dit_width": dit_width, + "timestep_dim": timestep_dim, + "feed_forward_dim": feed_forward_dim, + "mlp_hidden_dim": mlp_hidden_dim, + "state_dim": state_dim, + "action_dim": action_dim, + "future_token_count": future_token_count, + "max_sequence_length": max_sequence_length, + "block_count": PI_BLOCK_COUNT, + "tensor_count": len(expected), + "numel": sum(int(tensors[name].numel()) for name in expected), + } + + +def validate_pi_v3_tensors(tensors: dict[str, Any]) -> dict[str, int]: + """Validate every released PI_v3 policy tensor and infer its architecture.""" + actual = set(tensors) + expected = set(PI_V3_TENSOR_MAP) + missing = sorted(expected - actual) + if missing: + raise StarVLAError(f"PI-v3 policy is missing runtime tensors: {missing}") + + def shape(name: str) -> list[int]: + return [int(dim) for dim in tensors[name].shape] + + def matrix_shape(name: str) -> list[int]: + value = shape(name) + if len(value) != 2: + raise StarVLAError(f"invalid PI_v3 matrix shape for {name}: {value}") + return value + + timestep_input = matrix_shape( + "action_model.model.timestep_encoder.timestep_embedder.linear_1.weight" + ) + dit_width, timestep_dim = timestep_input + feed_forward_dim = matrix_shape( + "action_model.model.transformer_blocks.0.ff.net.0.proj.weight" + )[0] + mlp_hidden_dim = matrix_shape("action_model.action_decoder.layer1.weight")[0] + action_dim = matrix_shape("action_model.action_encoder.layer1.weight")[1] + future_token_count = matrix_shape("action_model.future_tokens.weight")[0] + max_sequence_length = matrix_shape("action_model.position_embedding.weight")[0] + qwen_hidden_dim = shape("project_layers.0.0.weight")[0] + projector_output_dim = matrix_shape("project_layers.0.1.weight")[0] + if projector_output_dim != dit_width: + raise StarVLAError( + "invalid PI_v3 projector/DiT width contract: " + f"projector={projector_output_dim}, DiT={dit_width}" + ) + + expected_shapes = { + "action_model.model.timestep_encoder.timestep_embedder.linear_1.weight": [ + dit_width, + timestep_dim, + ], + "action_model.model.timestep_encoder.timestep_embedder.linear_1.bias": [dit_width], + "action_model.model.timestep_encoder.timestep_embedder.linear_2.weight": [ + dit_width, + dit_width, + ], + "action_model.model.timestep_encoder.timestep_embedder.linear_2.bias": [dit_width], + "action_model.action_encoder.layer1.weight": [dit_width, action_dim], + "action_model.action_encoder.layer1.bias": [dit_width], + "action_model.action_encoder.layer2.weight": [dit_width, 2 * dit_width], + "action_model.action_encoder.layer2.bias": [dit_width], + "action_model.action_encoder.layer3.weight": [dit_width, dit_width], + "action_model.action_encoder.layer3.bias": [dit_width], + "action_model.action_decoder.layer1.weight": [mlp_hidden_dim, dit_width], + "action_model.action_decoder.layer1.bias": [mlp_hidden_dim], + "action_model.action_decoder.layer2.weight": [action_dim, mlp_hidden_dim], + "action_model.action_decoder.layer2.bias": [action_dim], + "action_model.future_tokens.weight": [future_token_count, dit_width], + "action_model.position_embedding.weight": [max_sequence_length, dit_width], + } + for block in range(PI_V3_BLOCK_COUNT): + prefix = f"action_model.model.transformer_blocks.{block}" + expected_shapes.update( + { + f"{prefix}.norm1.linear.weight": [2 * dit_width, dit_width], + f"{prefix}.norm1.linear.bias": [2 * dit_width], + f"{prefix}.attn1.to_q.weight": [dit_width, dit_width], + f"{prefix}.attn1.to_q.bias": [dit_width], + f"{prefix}.attn1.to_k.weight": [dit_width, dit_width], + f"{prefix}.attn1.to_k.bias": [dit_width], + f"{prefix}.attn1.to_v.weight": [dit_width, dit_width], + f"{prefix}.attn1.to_v.bias": [dit_width], + f"{prefix}.attn1.to_out.0.weight": [dit_width, dit_width], + f"{prefix}.attn1.to_out.0.bias": [dit_width], + f"{prefix}.ff.net.0.proj.weight": [feed_forward_dim, dit_width], + f"{prefix}.ff.net.0.proj.bias": [feed_forward_dim], + f"{prefix}.ff.net.2.weight": [dit_width, feed_forward_dim], + f"{prefix}.ff.net.2.bias": [dit_width], + } + ) + for projector in range(PI_V3_PROJECTOR_COUNT): + prefix = f"project_layers.{projector}" + expected_shapes.update( + { + f"{prefix}.0.weight": [qwen_hidden_dim], + f"{prefix}.0.bias": [qwen_hidden_dim], + f"{prefix}.1.weight": [dit_width, qwen_hidden_dim], + f"{prefix}.1.bias": [dit_width], + } + ) + mismatches = [ + f"{name}: expected {expected_shape}, got {shape(name)}" + for name, expected_shape in expected_shapes.items() + if shape(name) != expected_shape + ] + if mismatches: + raise StarVLAError("invalid PI_v3 tensor shapes: " + "; ".join(mismatches)) + + return { + "qwen_hidden_dim": qwen_hidden_dim, + "dit_width": dit_width, + "timestep_dim": timestep_dim, + "feed_forward_dim": feed_forward_dim, + "mlp_hidden_dim": mlp_hidden_dim, + "action_dim": action_dim, + "future_token_count": future_token_count, + "max_sequence_length": max_sequence_length, + "block_count": PI_V3_BLOCK_COUNT, + "projector_count": PI_V3_PROJECTOR_COUNT, + "tensor_count": len(PI_V3_TENSOR_MAP), + } + + +def load_variant_config( + policy_dir: Path, + surgery_manifest: dict[str, Any], + variant_name: str, +) -> dict[str, Any]: + catalog_variant = str(surgery_manifest.get("variant", variant_name)) + effective = resolve_effective_config( + policy_dir, + catalog_variant, + { + "framework": surgery_manifest.get("framework", variant_name), + "backbone": surgery_manifest.get("backbone", "qwen3_vl"), + }, + ) + effective_path = policy_dir / "effective_config.json" + effective_record = surgery_manifest.get("effective_config", {}) + if not effective_path.is_file(): + raise StarVLAError(f"missing surgery effective config: {effective_path}") + if ( + effective_record.get("path") != effective_path.name + or effective_record.get("size") != effective_path.stat().st_size + or effective_record.get("sha256") != sha256_file(effective_path) + ): + raise StarVLAError( + f"effective {variant_name.upper()} config does not match its canonical source/manifest" + ) + stored_effective = _load_json(effective_path) + if stored_effective != effective: + # Qwen3 bundles produced before Qwen2.5 support predate these two + # explicit annotations. Qwen3 was the only supported backbone then, so + # this is an unambiguous legacy spelling of the same effective config. + legacy_effective = copy.deepcopy(effective) + legacy_metadata = legacy_effective.get("_robotcpp_effective_config") + if ( + surgery_manifest.get("backbone", "qwen3_vl") == "qwen3_vl" + and isinstance(legacy_metadata, dict) + ): + legacy_metadata.pop("backbone", None) + legacy_metadata.pop("framework", None) + if stored_effective != legacy_effective: + raise StarVLAError( + f"effective {variant_name.upper()} config does not match its canonical source/manifest" + ) + return effective + + +def load_oft_config(policy_dir: Path, surgery_manifest: dict[str, Any]) -> dict[str, Any]: + return load_variant_config(policy_dir, surgery_manifest, "oft") + + +def load_groot_config(policy_dir: Path, surgery_manifest: dict[str, Any]) -> dict[str, Any]: + return load_variant_config(policy_dir, surgery_manifest, "groot") + + +def load_pi_config(policy_dir: Path, surgery_manifest: dict[str, Any]) -> dict[str, Any]: + return load_variant_config(policy_dir, surgery_manifest, "pi") + + +def load_pi_v3_config(policy_dir: Path, surgery_manifest: dict[str, Any]) -> dict[str, Any]: + return load_variant_config(policy_dir, surgery_manifest, "pi_v3") + + +def resolve_action_token_id(hf_dir: Path) -> int: + try: + from transformers import AutoTokenizer + except ImportError as exc: + raise StarVLAError("transformers is required to verify the OFT action token") from exc + try: + tokenizer = AutoTokenizer.from_pretrained(hf_dir, local_files_only=True, trust_remote_code=False) + token_ids = tokenizer(OFT_ACTION_TOKEN, add_special_tokens=False)["input_ids"] + except Exception as exc: + raise StarVLAError(f"failed to load the pinned Qwen tokenizer from {hf_dir}: {exc}") from exc + if token_ids != [OFT_ACTION_TOKEN_ID]: + raise StarVLAError( + f"unexpected OFT action token mapping for {OFT_ACTION_TOKEN!r}: " + f"expected [{OFT_ACTION_TOKEN_ID}], got {token_ids}" + ) + return token_ids[0] + + +def normalization_metadata(stats: dict[str, Any], action_dim: int) -> dict[str, Any]: + metadata: dict[str, Any] = { + "starvla.normalization.profile_count": len(stats), + "starvla.normalization.profile_keys": sorted(stats), + "starvla.normalization.clip_actions": False, + "starvla.normalization.binary_threshold": 0.5, + "starvla.normalization.binary_comparison": "gt", + } + for index, key in enumerate(sorted(stats)): + profile = stats[key] + action = profile.get("action") + if not isinstance(action, dict): + raise StarVLAError(f"normalization profile {key!r} has no action object") + for field in ("q01", "q99", "mask"): + values = action.get(field) + if not isinstance(values, list) or len(values) != action_dim: + raise StarVLAError( + f"normalization profile {key!r} action.{field} must have {action_dim} values" + ) + metadata[f"starvla.normalization.profile.{index}.action_{field}"] = values + q01 = action["q01"] + q99 = action["q99"] + mask = action["mask"] + expected_mask = [True] * (action_dim - 1) + [False] + if any(type(value) is not bool for value in mask) or mask != expected_mask: + raise StarVLAError( + f"normalization profile {key!r} action.mask must be {expected_mask}, got {mask}" + ) + if any( + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not math.isfinite(value) + for value in [*q01, *q99] + ): + raise StarVLAError(f"normalization profile {key!r} action quantiles must be finite numbers") + if any(q99[index] < q01[index] for index in range(action_dim - 1)): + raise StarVLAError(f"normalization profile {key!r} has q99 below q01") + metadata[f"starvla.normalization.profile.{index}.key"] = key + + state = profile.get("state") + if state is not None: + if not isinstance(state, dict): + raise StarVLAError(f"normalization profile {key!r} state must be an object") + state_q01 = state.get("q01") + state_q99 = state.get("q99") + if ( + not isinstance(state_q01, list) + or not isinstance(state_q99, list) + or not state_q01 + or len(state_q01) != len(state_q99) + ): + raise StarVLAError(f"normalization profile {key!r} has inconsistent state q01/q99") + if any( + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not math.isfinite(value) + for value in [*state_q01, *state_q99] + ): + raise StarVLAError(f"normalization profile {key!r} state quantiles must be finite numbers") + if any(upper < lower for lower, upper in zip(state_q01, state_q99)): + raise StarVLAError(f"normalization profile {key!r} state has q99 below q01") + metadata[f"starvla.normalization.profile.{index}.state_dimension"] = len(state_q01) + metadata[f"starvla.normalization.profile.{index}.state_q01"] = state_q01 + metadata[f"starvla.normalization.profile.{index}.state_q99"] = state_q99 + return metadata + + +def build_oft_metadata( + policy_dir: Path, + hf_dir: Path, + variant: dict[str, Any], + surgery_manifest: dict[str, Any], + dimensions: dict[str, int], + action_token_id: int, + text_filename: str, + mmproj_filename: str, +) -> dict[str, Any]: + backbone = str(variant.get("backbone", "qwen3_vl")) + config = load_oft_config(policy_dir, surgery_manifest) + framework = config.get("framework", {}) + action_config = framework.get("action_model", {}) + datasets = config.get("datasets", {}) + vla_config = datasets.get("vla_data", {}) + action_horizon = int(action_config.get("action_horizon", int(action_config.get("future_action_window_size", 15)) + 1)) + if action_horizon != 16: + raise StarVLAError(f"unexpected official OFT action horizon: {action_horizon}") + if dimensions["action_dim"] != 7: + raise StarVLAError(f"unexpected official OFT action dimension: {dimensions['action_dim']}") + expected_dimensions = { + "qwen3_vl": (2560, 5120), + "qwen2_5_vl": (2048, 4096), + }.get(backbone) + if expected_dimensions is None: + raise StarVLAError(f"unsupported OFT Qwen backbone: {backbone!r}") + if ( + dimensions["input_dim"], + dimensions["hidden_dim"], + ) != expected_dimensions: + raise StarVLAError(f"unexpected official OFT MLP dimensions: {dimensions}") + + action_tokens = OFT_ACTION_TOKEN * action_horizon + action_suffix = f" Please predict the next {action_horizon} robot actions: {action_tokens}." + image_size = vla_config.get("image_size", [224, 224]) + image_names = vla_config.get("obs", ["image_0"]) + if image_size != [224, 224] or image_names != ["image_0"]: + raise StarVLAError(f"unexpected official OFT image contract: image_size={image_size}, obs={image_names}") + + qwen = ( + _validate_pinned_qwen3vl_contract(hf_dir) + if backbone == "qwen3_vl" + else _validate_pinned_qwen25vl_contract(hf_dir) + ) + if qwen.get("hidden_size", dimensions["input_dim"]) != dimensions["input_dim"]: + raise StarVLAError( + "OFT policy input dimension does not match the staged Qwen backbone" + ) + + metadata: dict[str, Any] = { + "general.architecture": "starvla-policy", + "general.name": ( + "StarVLA Qwen3-VL OFT policy" + if backbone == "qwen3_vl" + else "StarVLA Qwen2.5-VL OFT policy" + ), + "starvla.schema_version": 1, + "starvla.framework": "oft", + "starvla.model_type": variant["model_type"], + "starvla.backbone.arch": backbone, + "starvla.bundle.uuid": surgery_manifest["bundle_uuid"], + "starvla.component.text.filename": text_filename, + "starvla.component.mmproj.filename": mmproj_filename, + "starvla.qwen.hidden_size": dimensions["input_dim"], + "starvla.qwen.input_embedding_size": ( + dimensions["input_dim"] * 4 + if backbone == "qwen3_vl" + else dimensions["input_dim"] + ), + "starvla.qwen.vocab_size": qwen.get("vocab_size", 151936), + "starvla.prompt.action_token": OFT_ACTION_TOKEN, + "starvla.prompt.action_token_id": action_token_id, + "starvla.prompt.action_suffix": action_suffix, + "starvla.prompt.cot_template": str(vla_config.get("CoT_prompt", "")), + "starvla.prompt.cot_enabled": bool(vla_config.get("CoT_prompt", "")), + "starvla.prompt.state_bins": 256, + "starvla.prompt.state_bin_min": -1.0, + "starvla.prompt.state_bin_max": 1.0, + "starvla.prompt.state_clip": False, + "starvla.action.dimension": dimensions["action_dim"], + "starvla.action.horizon": action_horizon, + "starvla.action.continuous_dimensions": [0, 1, 2, 3, 4, 5], + "starvla.action.binary_dimensions": [6], + "starvla.oft.hidden_size": dimensions["hidden_dim"], + "starvla.oft.block_count": 2, + "starvla.oft.layer_norm_epsilon": OFT_LAYER_NORM_EPS, + } + if backbone == "qwen3_vl": + metadata.update(_runtime_image_metadata( + build_qwen3vl_image_metadata( + vla_config, + qwen, + image_names, + variant_label="OFT", + ) + )) + else: + metadata.update(_runtime_image_metadata( + build_qwen25vl_image_metadata( + vla_config, + qwen, + image_names, + variant_label="OFT", + ) + )) + stats = _load_json(policy_dir / "dataset_statistics.json") + expected_profiles = ( + {"oxe_bridge", "oxe_rt1"} + if backbone == "qwen3_vl" + else {"bridge_dataset", "fractal20220817_data"} + ) + if set(stats) != expected_profiles: + raise StarVLAError(f"unexpected official OFT normalization profiles: {sorted(stats)}") + metadata.update(normalization_metadata(stats, dimensions["action_dim"])) + return metadata + + +def _validate_pinned_qwen3vl_contract(hf_dir: Path) -> dict[str, Any]: + qwen_config = _load_json(hf_dir / "config.json") + text_config = qwen_config.get("text_config", {}) + vision_config = qwen_config.get("vision_config", {}) + preprocessor = _load_json(hf_dir / "preprocessor_config.json") + actual = { + "architecture": qwen_config.get("architectures"), + "vocab_size": text_config.get("vocab_size"), + "hidden_size": text_config.get("hidden_size"), + "layer_count": text_config.get("num_hidden_layers"), + "head_count": text_config.get("num_attention_heads"), + "head_count_kv": text_config.get("num_key_value_heads"), + "head_dim": text_config.get("head_dim"), + "vision_hidden_size": vision_config.get("hidden_size"), + "vision_layer_count": vision_config.get("depth"), + "vision_head_count": vision_config.get("num_heads"), + "vision_patch_size": vision_config.get("patch_size"), + "vision_temporal_patch_size": vision_config.get("temporal_patch_size"), + "vision_merge_size": vision_config.get("spatial_merge_size"), + "vision_deepstack": vision_config.get("deepstack_visual_indexes"), + "processor_size": preprocessor.get("size"), + "processor_patch_size": preprocessor.get("patch_size"), + "processor_temporal_patch_size": preprocessor.get("temporal_patch_size"), + "processor_merge_size": preprocessor.get("merge_size"), + "processor_class": preprocessor.get("processor_class"), + "image_processor_type": preprocessor.get("image_processor_type"), + "image_mean": preprocessor.get("image_mean"), + "image_std": preprocessor.get("image_std"), + } + expected = { + "architecture": ["Qwen3VLForConditionalGeneration"], + "vocab_size": 151936, + "hidden_size": 2560, + "layer_count": 36, + "head_count": 32, + "head_count_kv": 8, + "head_dim": 128, + "vision_hidden_size": 1024, + "vision_layer_count": 24, + "vision_head_count": 16, + "vision_patch_size": 16, + "vision_temporal_patch_size": 2, + "vision_merge_size": 2, + "vision_deepstack": [5, 11, 17], + "processor_size": { + "shortest_edge": QWEN3VL_PROCESSOR_MIN_PIXELS, + "longest_edge": QWEN3VL_PROCESSOR_MAX_PIXELS, + }, + "processor_patch_size": QWEN3VL_IMAGE_PATCH_SIZE, + "processor_temporal_patch_size": QWEN3VL_TEMPORAL_PATCH_SIZE, + "processor_merge_size": QWEN3VL_SPATIAL_MERGE_SIZE, + "processor_class": "Qwen3VLProcessor", + "image_processor_type": "Qwen2VLImageProcessorFast", + "image_mean": QWEN3VL_IMAGE_MEAN, + "image_std": QWEN3VL_IMAGE_STD, + } + if actual != expected: + raise StarVLAError(f"unexpected pinned Qwen config/processor contract: {actual}") + chat_template_path = hf_dir / "chat_template.json" + if not chat_template_path.is_file(): + raise StarVLAError(f"missing pinned Qwen chat template: {chat_template_path}") + return { + **actual, + "chat_template_sha256": sha256_file(chat_template_path), + } + + +def _validate_pinned_qwen25vl_contract(hf_dir: Path) -> dict[str, Any]: + qwen_config = _load_json(hf_dir / "config.json") + text_config = qwen_config.get("text_config") + if not isinstance(text_config, dict): + text_config = qwen_config + vision_config = qwen_config.get("vision_config", {}) + preprocessor = _load_json(hf_dir / "preprocessor_config.json") + hidden_size = text_config.get("hidden_size") + head_count = text_config.get("num_attention_heads") + head_dim = ( + hidden_size // head_count + if isinstance(hidden_size, int) + and isinstance(head_count, int) + and head_count > 0 + and hidden_size % head_count == 0 + else None + ) + actual = { + "architecture": qwen_config.get("architectures"), + "model_type": qwen_config.get("model_type"), + "tie_word_embeddings": text_config.get("tie_word_embeddings"), + "vocab_size": text_config.get("vocab_size"), + "hidden_size": hidden_size, + "layer_count": text_config.get("num_hidden_layers"), + "head_count": head_count, + "head_count_kv": text_config.get("num_key_value_heads"), + "head_dim": head_dim, + "vision_hidden_size": vision_config.get("hidden_size"), + "vision_layer_count": vision_config.get("depth"), + "vision_head_count": vision_config.get("num_heads"), + "vision_patch_size": vision_config.get("patch_size"), + "vision_temporal_patch_size": vision_config.get("temporal_patch_size"), + "vision_merge_size": vision_config.get("spatial_merge_size"), + "vision_window_size": vision_config.get("window_size"), + "vision_full_attention_blocks": vision_config.get("fullatt_block_indexes"), + "vision_deepstack": [], + "processor_min_pixels": preprocessor.get("min_pixels"), + "processor_max_pixels": preprocessor.get("max_pixels"), + "processor_patch_size": preprocessor.get("patch_size"), + "processor_temporal_patch_size": preprocessor.get("temporal_patch_size"), + "processor_merge_size": preprocessor.get("merge_size"), + "processor_class": preprocessor.get("processor_class"), + "image_processor_type": preprocessor.get("image_processor_type"), + "image_mean": preprocessor.get("image_mean"), + "image_std": preprocessor.get("image_std"), + } + expected_image_processor_type = { + 151_936: "Qwen2VLImageProcessor", + 153_713: "Qwen2VLImageProcessorFast", + }.get(actual["vocab_size"]) + expected = { + "architecture": ["Qwen2_5_VLForConditionalGeneration"], + "model_type": "qwen2_5_vl", + "tie_word_embeddings": False, + "vocab_size": actual["vocab_size"], + "hidden_size": 2048, + "layer_count": 36, + "head_count": 16, + "head_count_kv": 2, + "head_dim": 128, + "vision_hidden_size": 1280, + "vision_layer_count": 32, + "vision_head_count": 16, + "vision_patch_size": QWEN25VL_IMAGE_PATCH_SIZE, + "vision_temporal_patch_size": QWEN25VL_TEMPORAL_PATCH_SIZE, + "vision_merge_size": QWEN25VL_SPATIAL_MERGE_SIZE, + "vision_window_size": 112, + "vision_full_attention_blocks": [7, 15, 23, 31], + "vision_deepstack": [], + "processor_min_pixels": QWEN25VL_PROCESSOR_MIN_PIXELS, + "processor_max_pixels": QWEN25VL_PROCESSOR_MAX_PIXELS, + "processor_patch_size": QWEN25VL_IMAGE_PATCH_SIZE, + "processor_temporal_patch_size": QWEN25VL_TEMPORAL_PATCH_SIZE, + "processor_merge_size": QWEN25VL_SPATIAL_MERGE_SIZE, + "processor_class": "Qwen2_5_VLProcessor", + "image_processor_type": expected_image_processor_type, + "image_mean": QWEN25VL_IMAGE_MEAN, + "image_std": QWEN25VL_IMAGE_STD, + } + if expected_image_processor_type is None: + raise StarVLAError( + f"unexpected pinned Qwen2.5-VL vocabulary size: {actual['vocab_size']!r}" + ) + if actual != expected: + raise StarVLAError( + f"unexpected pinned Qwen2.5-VL config/processor contract: {actual}" + ) + chat_template_path = hf_dir / "chat_template.json" + if not chat_template_path.is_file(): + # The action-expanded checkpoint publishes the same template as Jinja. + chat_template_path = hf_dir / "chat_template.jinja" + if not chat_template_path.is_file(): + raise StarVLAError(f"missing pinned Qwen2.5-VL chat template in {hf_dir}") + return { + **actual, + "chat_template_sha256": sha256_file(chat_template_path), + } + + +def _validate_pinned_qwenvl_contract( + hf_dir: Path, backbone: str +) -> dict[str, Any]: + if backbone == "qwen3_vl": + return _validate_pinned_qwen3vl_contract(hf_dir) + if backbone == "qwen2_5_vl": + return _validate_pinned_qwen25vl_contract(hf_dir) + raise StarVLAError(f"unsupported StarVLA Qwen backbone: {backbone!r}") + + +def _require_released_obs_pre_resize_disabled( + vla_config: dict[str, Any], + variant_label: str, + config_label: str, +) -> None: + if not isinstance(vla_config, dict): + raise StarVLAError(f"official {variant_label} {config_label} vla_data must be an object") + if "obs_image_size" in vla_config: + raise StarVLAError( + f"official {variant_label} {config_label} unexpectedly defines " + "datasets.vla_data.obs_image_size; released predict_action must leave its " + "optional pre-resize branch disabled" + ) + + +def build_qwen3vl_image_metadata( + vla_config: dict[str, Any], + qwen: dict[str, Any], + image_names: list[str], + *, + variant_label: str, + config_label: str = "effective config", +) -> dict[str, Any]: + """Build the released dynamic Qwen3-VL image preprocessing contract.""" + _require_released_obs_pre_resize_disabled(vla_config, variant_label, config_label) + if image_names != ["image_0"]: + raise StarVLAError(f"unexpected official {variant_label} image names: {image_names!r}") + + processor_size = qwen.get("processor_size") + actual = { + "min_pixels": processor_size.get("shortest_edge") if isinstance(processor_size, dict) else None, + "max_pixels": processor_size.get("longest_edge") if isinstance(processor_size, dict) else None, + "processor_patch_size": qwen.get("processor_patch_size"), + "processor_temporal_patch_size": qwen.get("processor_temporal_patch_size"), + "processor_merge_size": qwen.get("processor_merge_size"), + "processor_class": qwen.get("image_processor_type"), + "image_mean": qwen.get("image_mean"), + "image_std": qwen.get("image_std"), + "vision_patch_size": qwen.get("vision_patch_size"), + "vision_temporal_patch_size": qwen.get("vision_temporal_patch_size"), + "vision_merge_size": qwen.get("vision_merge_size"), + } + expected = { + "min_pixels": QWEN3VL_PROCESSOR_MIN_PIXELS, + "max_pixels": QWEN3VL_PROCESSOR_MAX_PIXELS, + "processor_patch_size": QWEN3VL_IMAGE_PATCH_SIZE, + "processor_temporal_patch_size": QWEN3VL_TEMPORAL_PATCH_SIZE, + "processor_merge_size": QWEN3VL_SPATIAL_MERGE_SIZE, + "processor_class": "Qwen2VLImageProcessorFast", + "image_mean": QWEN3VL_IMAGE_MEAN, + "image_std": QWEN3VL_IMAGE_STD, + "vision_patch_size": QWEN3VL_IMAGE_PATCH_SIZE, + "vision_temporal_patch_size": QWEN3VL_TEMPORAL_PATCH_SIZE, + "vision_merge_size": QWEN3VL_SPATIAL_MERGE_SIZE, + } + if actual != expected: + raise StarVLAError(f"unexpected pinned Qwen dynamic image contract: {actual}") + + token_area = QWEN3VL_IMAGE_PATCH_SIZE**2 * QWEN3VL_SPATIAL_MERGE_SIZE**2 + if ( + QWEN3VL_PROCESSOR_MIN_PIXELS // token_area != QWEN3VL_MIN_IMAGE_TOKENS + or QWEN3VL_PROCESSOR_MAX_PIXELS // token_area != QWEN3VL_MAX_IMAGE_TOKENS + or QWEN3VL_PROCESSOR_MIN_PIXELS % token_area + or QWEN3VL_PROCESSOR_MAX_PIXELS % token_area + ): + raise StarVLAError("internal Qwen3-VL smart-resize image-token bounds drift") + return { + key: list(value) if isinstance(value, list) else value + for key, value in QWEN3VL_DYNAMIC_IMAGE_METADATA.items() + } + + +def build_qwen25vl_image_metadata( + vla_config: dict[str, Any], + qwen: dict[str, Any], + image_names: list[str], + *, + variant_label: str, + config_label: str = "effective config", +) -> dict[str, Any]: + """Build the Transformers 4.57 fast Qwen2.5-VL image contract.""" + _require_released_obs_pre_resize_disabled( + vla_config, variant_label, config_label + ) + if image_names != ["image_0"]: + raise StarVLAError( + f"unexpected official {variant_label} image names: {image_names!r}" + ) + expected = { + "processor_min_pixels": QWEN25VL_PROCESSOR_MIN_PIXELS, + "processor_max_pixels": QWEN25VL_PROCESSOR_MAX_PIXELS, + "processor_patch_size": QWEN25VL_IMAGE_PATCH_SIZE, + "processor_temporal_patch_size": QWEN25VL_TEMPORAL_PATCH_SIZE, + "processor_merge_size": QWEN25VL_SPATIAL_MERGE_SIZE, + "image_mean": QWEN25VL_IMAGE_MEAN, + "image_std": QWEN25VL_IMAGE_STD, + "vision_patch_size": QWEN25VL_IMAGE_PATCH_SIZE, + "vision_temporal_patch_size": QWEN25VL_TEMPORAL_PATCH_SIZE, + "vision_merge_size": QWEN25VL_SPATIAL_MERGE_SIZE, + } + actual = {key: qwen.get(key) for key in expected} + if actual != expected: + raise StarVLAError( + f"unexpected pinned Qwen2.5-VL dynamic image contract: {actual}" + ) + + token_area = ( + QWEN25VL_IMAGE_PATCH_SIZE**2 * QWEN25VL_SPATIAL_MERGE_SIZE**2 + ) + if ( + QWEN25VL_PROCESSOR_MIN_PIXELS // token_area + != QWEN25VL_MIN_IMAGE_TOKENS + or QWEN25VL_PROCESSOR_MAX_PIXELS // token_area + != QWEN25VL_MAX_IMAGE_TOKENS + or QWEN25VL_PROCESSOR_MIN_PIXELS % token_area + or QWEN25VL_PROCESSOR_MAX_PIXELS % token_area + ): + raise StarVLAError( + "internal Qwen2.5-VL smart-resize image-token bounds drift" + ) + return { + "starvla.image.count": 1, + "starvla.image.names": list(image_names), + "starvla.image.preprocessing_mode": "qwen2_5vl_smart_resize", + "starvla.image.framework_inference_pre_resize": False, + "starvla.image.framework_inference_pre_resize_config_key": + "datasets.vla_data.obs_image_size", + "starvla.image.processor_min_pixels": QWEN25VL_PROCESSOR_MIN_PIXELS, + "starvla.image.processor_max_pixels": QWEN25VL_PROCESSOR_MAX_PIXELS, + "starvla.image.processor_class": "Qwen2VLImageProcessorFast", + "starvla.image.processor_reference_transformers_version": "4.57.0", + "starvla.image.processor_do_convert_rgb": True, + "starvla.image.processor_do_resize": True, + "starvla.image.processor_resize_resample": "bicubic", + "starvla.image.processor_resize_antialias": True, + "starvla.image.processor_do_rescale": True, + "starvla.image.processor_rescale_factor": 1.0 / 255.0, + "starvla.image.processor_do_normalize": True, + "starvla.image.processor_image_mean": list(QWEN25VL_IMAGE_MEAN), + "starvla.image.processor_image_std": list(QWEN25VL_IMAGE_STD), + "starvla.image.patch_size": QWEN25VL_IMAGE_PATCH_SIZE, + "starvla.image.temporal_patch_size": QWEN25VL_TEMPORAL_PATCH_SIZE, + "starvla.image.spatial_merge_size": QWEN25VL_SPATIAL_MERGE_SIZE, + "starvla.image.token_count_mode": + "dynamic_grid_thw_after_spatial_merge", + "starvla.image.min_token_count": QWEN25VL_MIN_IMAGE_TOKENS, + "starvla.image.max_token_count": QWEN25VL_MAX_IMAGE_TOKENS, + } + + +def _runtime_image_metadata(metadata: dict[str, Any]) -> dict[str, Any]: + keys = ( + "starvla.image.count", + "starvla.image.names", + "starvla.image.processor_min_pixels", + "starvla.image.processor_max_pixels", + "starvla.image.patch_size", + "starvla.image.spatial_merge_size", + "starvla.image.min_token_count", + "starvla.image.max_token_count", + ) + return {key: metadata[key] for key in keys} + + +def build_groot_metadata( + policy_dir: Path, + hf_dir: Path, + variant: dict[str, Any], + surgery_manifest: dict[str, Any], + dimensions: dict[str, int], + text_filename: str, + mmproj_filename: str, +) -> dict[str, Any]: + """Build the executable contract for a released Qwen-VL GR00T head.""" + backbone = str(variant.get("backbone", "qwen3_vl")) + config = load_groot_config(policy_dir, surgery_manifest) + framework = config.get("framework", {}) + action_config = framework.get("action_model", {}) + diffusion_config = action_config.get("diffusion_model_cfg", {}) + vla_config = config.get("datasets", {}).get("vla_data", {}) + qwen = _validate_pinned_qwenvl_contract(hf_dir, backbone) + + expected_dimensions = GROOT_OFFICIAL_DIMENSIONS_BY_BACKBONE.get(backbone) + if expected_dimensions is None: + raise StarVLAError(f"unsupported GR00T Qwen backbone: {backbone!r}") + if dimensions != expected_dimensions: + raise StarVLAError(f"unexpected official GR00T tensor dimensions: {dimensions}") + if qwen["hidden_size"] != dimensions["qwen_hidden_dim"]: + raise StarVLAError( + "GR00T cross-attention dimension does not match the staged Qwen backbone" + ) + + action_horizon = int( + action_config.get( + "action_horizon", + int(action_config.get("future_action_window_size", 15)) + 1, + ) + ) + expected_action_config = { + "action_model_type": action_config.get("action_model_type"), + "hidden_size": action_config.get("hidden_size"), + "add_pos_embed": action_config.get("add_pos_embed"), + "max_seq_len": action_config.get("max_seq_len"), + "action_dim": action_config.get("action_dim"), + "state_dim": action_config.get("state_dim"), + "action_horizon": action_horizon, + "past_action_window_size": action_config.get("past_action_window_size"), + "repeated_diffusion_steps": action_config.get("repeated_diffusion_steps"), + "noise_beta_alpha": action_config.get("noise_beta_alpha"), + "noise_beta_beta": action_config.get("noise_beta_beta"), + "noise_s": action_config.get("noise_s"), + "num_timestep_buckets": action_config.get("num_timestep_buckets"), + "num_inference_timesteps": action_config.get("num_inference_timesteps"), + "num_target_vision_tokens": action_config.get("num_target_vision_tokens"), + } + required_action_config = { + "action_model_type": "DiT-B", + "hidden_size": 1024, + "add_pos_embed": True, + "max_seq_len": 1024, + "action_dim": 7, + "state_dim": 7, + "action_horizon": 16, + "past_action_window_size": 0, + "repeated_diffusion_steps": 8, + "noise_beta_alpha": 1.5, + "noise_beta_beta": 1.0, + "noise_s": 0.999, + "num_timestep_buckets": 1000, + "num_inference_timesteps": 4, + "num_target_vision_tokens": 32, + } + if expected_action_config != required_action_config: + raise StarVLAError(f"unexpected official GR00T action config: {expected_action_config}") + + actual_diffusion_config = { + "input_embedding_dim": diffusion_config.get("input_embedding_dim"), + "attention_head_dim": diffusion_config.get("attention_head_dim"), + "num_attention_heads": diffusion_config.get("num_attention_heads"), + "cross_attention_dim": diffusion_config.get("cross_attention_dim"), + "dropout": diffusion_config.get("dropout"), + "final_dropout": diffusion_config.get("final_dropout"), + "interleave_self_attention": diffusion_config.get("interleave_self_attention"), + "norm_type": diffusion_config.get("norm_type"), + "num_layers": diffusion_config.get("num_layers"), + "output_dim": diffusion_config.get("output_dim"), + "positional_embeddings": diffusion_config.get("positional_embeddings"), + } + required_diffusion_config = { + "input_embedding_dim": 768, + "attention_head_dim": 64, + "num_attention_heads": 12, + "cross_attention_dim": dimensions["qwen_hidden_dim"], + "dropout": 0.2, + "final_dropout": True, + "interleave_self_attention": True, + "norm_type": "ada_norm", + "num_layers": 16, + "output_dim": 1024, + "positional_embeddings": None, + } + if actual_diffusion_config != required_diffusion_config: + raise StarVLAError(f"unexpected official GR00T diffusion config: {actual_diffusion_config}") + + framework_identity = ( + framework.get("name") + if backbone == "qwen3_vl" + else framework.get("framework_py") + ) + expected_framework_identity = ( + "QwenGR00T" if backbone == "qwen3_vl" else "QwenFM" + ) + if framework_identity != expected_framework_identity: + raise StarVLAError( + f"unexpected official GR00T framework identity: {framework_identity!r}" + ) + if vla_config.get("image_size") != [224, 224] or vla_config.get("obs") != ["image_0"]: + raise StarVLAError( + "unexpected official GR00T image contract: " + f"image_size={vla_config.get('image_size')}, obs={vla_config.get('obs')}" + ) + if vla_config.get("include_state", False) not in (False, "False"): + raise StarVLAError("released GR00T checkpoint unexpectedly enables training state input") + + cot_template = str(vla_config.get("CoT_prompt", "")) + required_cot = ( + "Your task is {instruction}. To identify the key objects for your task. " + "Locate their bounding boxes in [x1,y1,x2,y2] format." + ) + if cot_template != required_cot: + raise StarVLAError(f"unexpected official GR00T CoT prompt: {cot_template!r}") + + num_steps = int(action_config["num_inference_timesteps"]) + timestep_buckets = int(action_config["num_timestep_buckets"]) + timestep_ids = [step * timestep_buckets // num_steps for step in range(num_steps)] + metadata: dict[str, Any] = { + "general.architecture": "starvla-policy", + "general.name": ( + "StarVLA Qwen3-VL GR00T policy" + if backbone == "qwen3_vl" + else "StarVLA Qwen2.5-VL GR00T policy" + ), + "starvla.schema_version": 1, + "starvla.framework": "groot", + "starvla.model_type": variant["model_type"], + "starvla.backbone.arch": backbone, + "starvla.bundle.uuid": surgery_manifest["bundle_uuid"], + "starvla.component.text.filename": text_filename, + "starvla.component.mmproj.filename": mmproj_filename, + "starvla.qwen.hidden_size": dimensions["qwen_hidden_dim"], + "starvla.qwen.input_embedding_size": ( + dimensions["qwen_hidden_dim"] * 4 + if backbone == "qwen3_vl" + else dimensions["qwen_hidden_dim"] + ), + "starvla.qwen.vocab_size": qwen["vocab_size"], + "starvla.prompt.cot_template": cot_template, + "starvla.action.dimension": dimensions["action_dim"], + "starvla.action.horizon": action_horizon, + "starvla.action.continuous_dimensions": [0, 1, 2, 3, 4, 5], + "starvla.action.binary_dimensions": [6], + "starvla.groot.dit_width": dimensions["dit_width"], + "starvla.groot.block_count": dimensions["block_count"], + "starvla.groot.attention_head_count": 12, + "starvla.groot.attention_head_dim": 64, + "starvla.groot.cross_attention_dim": dimensions["qwen_hidden_dim"], + "starvla.groot.feed_forward_dim": dimensions["feed_forward_dim"], + "starvla.groot.ada_norm_epsilon": GROOT_DIT_NORM_EPS, + "starvla.groot.output_norm_epsilon": GROOT_OUTPUT_NORM_EPS, + "starvla.groot.output_dimension": dimensions["output_dim"], + "starvla.groot.mlp_hidden_dimension": dimensions["mlp_hidden_dim"], + "starvla.groot.future_token_count": dimensions["future_token_count"], + "starvla.groot.action_position_count": dimensions["max_sequence_length"], + "starvla.groot.no_state_sequence_length": dimensions["future_token_count"] + action_horizon, + "starvla.groot.timestep_projection_dim": dimensions["timestep_dim"], + "starvla.groot.timestep_ids": timestep_ids, + "starvla.groot.euler_dt": 1.0 / num_steps, + } + if backbone == "qwen3_vl": + metadata.update(_runtime_image_metadata( + build_qwen3vl_image_metadata( + vla_config, + qwen, + ["image_0"], + variant_label="GR00T", + ) + )) + else: + metadata.update(_runtime_image_metadata( + build_qwen25vl_image_metadata( + vla_config, + qwen, + ["image_0"], + variant_label="GR00T", + ) + )) + stats = _load_json(policy_dir / "dataset_statistics.json") + if set(stats) != {"oxe_bridge", "oxe_rt1"}: + raise StarVLAError(f"unexpected official GR00T normalization profiles: {sorted(stats)}") + state_dimensions = sorted( + { + len(profile.get("state", {}).get("q01", [])) + for profile in stats.values() + } + ) + if state_dimensions != [8]: + raise StarVLAError(f"unexpected official GR00T state statistics dimensions: {state_dimensions}") + metadata.update(normalization_metadata(stats, dimensions["action_dim"])) + return metadata + + +def build_pi_metadata( + policy_dir: Path, + hf_dir: Path, + variant: dict[str, Any], + surgery_manifest: dict[str, Any], + dimensions: dict[str, int], + text_filename: str, + mmproj_filename: str, +) -> dict[str, Any]: + """Build the released Qwen2.5-VL legacy PI executable contract.""" + if variant.get("framework") != "pi" or variant.get("backbone") != "qwen2_5_vl": + raise StarVLAError("legacy PI metadata requires the qwen25_pi catalog variant") + config = load_pi_config(policy_dir, surgery_manifest) + framework = config.get("framework", {}) + action_config = framework.get("action_model", {}) + diffusion_config = action_config.get("diffusion_model_cfg", {}) + vla_config = config.get("datasets", {}).get("vla_data", {}) + qwen = _validate_pinned_qwen25vl_contract(hf_dir) + + if dimensions != PI_OFFICIAL_DIMENSIONS: + raise StarVLAError(f"unexpected official legacy PI tensor dimensions: {dimensions}") + if qwen["hidden_size"] != dimensions["qwen_hidden_dim"]: + raise StarVLAError( + "legacy PI cross-attention dimension does not match the staged Qwen backbone" + ) + + actual_action_config = { + "action_model_type": action_config.get("action_model_type"), + "hidden_size": action_config.get("hidden_size"), + "action_hidden_dim": action_config.get("action_hidden_dim"), + "add_pos_embed": action_config.get("add_pos_embed"), + "max_seq_len": action_config.get("max_seq_len"), + "action_dim": action_config.get("action_dim"), + "state_dim": action_config.get("state_dim"), + "future_action_window_size": action_config.get("future_action_window_size"), + "action_horizon": action_config.get("action_horizon"), + "past_action_window_size": action_config.get("past_action_window_size"), + "repeated_diffusion_steps": action_config.get("repeated_diffusion_steps"), + "noise_beta_alpha": action_config.get("noise_beta_alpha"), + "noise_beta_beta": action_config.get("noise_beta_beta"), + "noise_s": action_config.get("noise_s"), + "num_timestep_buckets": action_config.get("num_timestep_buckets"), + "num_inference_timesteps": action_config.get("num_inference_timesteps"), + "num_target_vision_tokens": action_config.get("num_target_vision_tokens"), + } + required_action_config = { + "action_model_type": "DiT-Qwen", + "hidden_size": 2048, + "action_hidden_dim": 2048, + "add_pos_embed": True, + "max_seq_len": 1024, + "action_dim": 7, + "state_dim": 7, + "future_action_window_size": 15, + "action_horizon": 16, + "past_action_window_size": 0, + "repeated_diffusion_steps": 8, + "noise_beta_alpha": 1.5, + "noise_beta_beta": 1.0, + "noise_s": 0.999, + "num_timestep_buckets": 1000, + "num_inference_timesteps": 4, + "num_target_vision_tokens": 32, + } + if actual_action_config != required_action_config: + raise StarVLAError( + f"unexpected official legacy PI action config: {actual_action_config}" + ) + + actual_diffusion_config = { + "input_embedding_dim": diffusion_config.get("input_embedding_dim"), + "attention_head_dim": diffusion_config.get("attention_head_dim"), + "num_attention_heads": diffusion_config.get("num_attention_heads"), + "cross_attention_dim": diffusion_config.get("cross_attention_dim"), + "dropout": diffusion_config.get("dropout"), + "final_dropout": diffusion_config.get("final_dropout"), + "interleave_self_attention": diffusion_config.get("interleave_self_attention"), + "use_canonical_forward": diffusion_config.get("use_canonical_forward"), + "norm_type": diffusion_config.get("norm_type"), + "num_layers": diffusion_config.get("num_layers"), + "output_dim": diffusion_config.get("output_dim"), + "positional_embeddings": diffusion_config.get("positional_embeddings"), + } + required_diffusion_config = { + "input_embedding_dim": 2048, + "attention_head_dim": 64, + "num_attention_heads": 32, + "cross_attention_dim": 2048, + "dropout": 0.2, + "final_dropout": True, + "interleave_self_attention": True, + "use_canonical_forward": False, + "norm_type": "ada_norm", + "num_layers": 16, + "output_dim": 1024, + "positional_embeddings": None, + } + if actual_diffusion_config != required_diffusion_config: + raise StarVLAError( + f"unexpected official legacy PI diffusion config: {actual_diffusion_config}" + ) + if framework.get("name") != "QwenPI": + raise StarVLAError( + f"unexpected official legacy PI framework name: {framework.get('name')!r}" + ) + qwen_config = framework.get("qwenvl", {}) + if ( + qwen_config.get("vl_hidden_dim") != dimensions["qwen_hidden_dim"] + or qwen_config.get("attn_implementation") != "flash_attention_2" + ): + raise StarVLAError( + f"unexpected official legacy PI Qwen contract: {qwen_config}" + ) + + required_cot = ( + "Your task is {instruction}. To identify the key objects for your task. " + "Locate their bounding boxes in [x1,y1,x2,y2] format." + ) + cot_template = str(vla_config.get("CoT_prompt", "")) + if ( + cot_template != required_cot + or vla_config.get("obs") != ["image_0"] + or vla_config.get("image_size") != [224, 224] + or vla_config.get("default_image_resolution") != [3, 224, 224] + or vla_config.get("data_mix") != "bridge_rt_1" + or vla_config.get("action_type") != "delta_ee" + ): + raise StarVLAError(f"unexpected official legacy PI VLA config: {vla_config}") + + num_steps = int(action_config["num_inference_timesteps"]) + timestep_buckets = int(action_config["num_timestep_buckets"]) + continuous_times = [step / float(num_steps) for step in range(num_steps)] + timestep_ids = [int(value * timestep_buckets) for value in continuous_times] + hidden_tuple_indices = list( + range(qwen["layer_count"] - PI_BLOCK_COUNT + 1, qwen["layer_count"] + 1) + ) + metadata: dict[str, Any] = { + "general.architecture": "starvla-policy", + "general.name": "StarVLA Qwen2.5-VL legacy PI policy", + "general.source.uuid": surgery_manifest["bundle_uuid"], + "starvla.schema_version": 1, + "starvla.framework": "pi", + "starvla.model_type": variant["model_type"], + "starvla.backbone.arch": "qwen2_5_vl", + "starvla.bundle.uuid": surgery_manifest["bundle_uuid"], + "starvla.component.text.filename": text_filename, + "starvla.component.mmproj.filename": mmproj_filename, + "starvla.qwen.hidden_size": dimensions["qwen_hidden_dim"], + "starvla.qwen.input_embedding_size": dimensions["qwen_hidden_dim"], + "starvla.qwen.layer_count": qwen["layer_count"], + "starvla.qwen.vocab_size": qwen["vocab_size"], + "starvla.prompt.cot_template": cot_template, + "starvla.conditioning.hidden_tuple_indices": hidden_tuple_indices, + "starvla.action.dimension": dimensions["action_dim"], + "starvla.action.horizon": 16, + "starvla.action.continuous_dimensions": [0, 1, 2, 3, 4, 5], + "starvla.action.binary_dimensions": [6], + "starvla.state.dimension": dimensions["state_dim"], + "starvla.pi.dit_width": dimensions["dit_width"], + "starvla.pi.block_count": dimensions["block_count"], + "starvla.pi.attention_head_count": 32, + "starvla.pi.attention_head_dim": 64, + "starvla.pi.cross_attention_dim": dimensions["qwen_hidden_dim"], + "starvla.pi.feed_forward_dim": dimensions["feed_forward_dim"], + "starvla.pi.mlp_hidden_dimension": dimensions["mlp_hidden_dim"], + "starvla.pi.state_token_count": 1, + "starvla.pi.future_token_count": dimensions["future_token_count"], + "starvla.pi.action_position_count": dimensions["max_sequence_length"], + "starvla.pi.timestep_projection_dim": dimensions["timestep_dim"], + "starvla.pi.num_inference_timesteps": num_steps, + "starvla.pi.timestep_ids": timestep_ids, + "starvla.pi.euler_dt": 1.0 / num_steps, + "starvla.pi.ada_norm_epsilon": PI_DIT_NORM_EPS, + } + image_metadata = build_qwen25vl_image_metadata( + vla_config, + qwen, + ["image_0"], + variant_label="legacy PI", + ) + image_metadata.update( + { + "starvla.image.framework_inference_pre_resize": True, + "starvla.image.framework_inference_pre_resize_config_key": + "datasets.vla_data.image_size", + "starvla.image.framework_inference_pre_resize_width": 224, + "starvla.image.framework_inference_pre_resize_height": 224, + } + ) + for key in ( + "starvla.image.count", + "starvla.image.names", + "starvla.image.processor_min_pixels", + "starvla.image.processor_max_pixels", + "starvla.image.patch_size", + "starvla.image.spatial_merge_size", + "starvla.image.min_token_count", + "starvla.image.max_token_count", + "starvla.image.framework_inference_pre_resize_width", + "starvla.image.framework_inference_pre_resize_height", + ): + metadata[key] = image_metadata[key] + + stats = _load_json(policy_dir / "dataset_statistics.json") + if set(stats) != {"oxe_bridge", "oxe_rt1"}: + raise StarVLAError( + f"unexpected official legacy PI normalization profiles: {sorted(stats)}" + ) + state_dimensions = sorted( + { + len(profile.get("state", {}).get("q01", [])) + for profile in stats.values() + } + ) + if state_dimensions != [8]: + raise StarVLAError( + f"unexpected official legacy PI state statistics dimensions: {state_dimensions}" + ) + metadata.update(normalization_metadata(stats, dimensions["action_dim"])) + metadata["starvla.normalization.clip_actions"] = True + metadata["starvla.normalization.binary_comparison"] = "ge" + return metadata + + +def build_pi_v3_metadata( + policy_dir: Path, + hf_dir: Path, + variant: dict[str, Any], + surgery_manifest: dict[str, Any], + dimensions: dict[str, int], + text_filename: str, + mmproj_filename: str, +) -> dict[str, Any]: + config = load_pi_v3_config(policy_dir, surgery_manifest) + full_config = _load_yaml(policy_dir / "config.full.yaml") + framework = config.get("framework", {}) + action = framework.get("action_model", {}) + diffusion = action.get("diffusion_model_cfg", {}) + vla = config.get("datasets", {}).get("vla_data", {}) + image_names = full_config.get("datasets", {}).get("vla_data", {}).get("obs") + qwen = _validate_pinned_qwen3vl_contract(hf_dir) + + expected_dimensions = { + "qwen_hidden_dim": qwen.get("hidden_size"), + "dit_width": diffusion.get("action_dit_hidden_dim"), + "action_dim": action.get("action_dim"), + "block_count": diffusion.get("num_layers"), + } + if framework.get("name") != "QwenPI_v3" or any( + dimensions[key] != value for key, value in expected_dimensions.items() + ): + raise StarVLAError("PI-v3 config does not match the checkpoint tensor shapes") + if not isinstance(image_names, list) or not image_names: + raise StarVLAError("PI-v3 config does not define observation image names") + + horizon = int(action["action_horizon"]) + num_steps = int(action["num_inference_timesteps"]) + timestep_buckets = int(action["num_timestep_buckets"]) + processor_size = qwen["processor_size"] + metadata: dict[str, Any] = { + "general.architecture": "starvla-policy", + "general.name": "StarVLA Qwen3-VL PI-v3 policy", + "general.source.uuid": surgery_manifest["bundle_uuid"], + "starvla.schema_version": 1, + "starvla.framework": "pi_v3", + "starvla.model_type": "starvla", + "starvla.backbone.arch": "qwen3_vl", + "starvla.bundle.uuid": surgery_manifest["bundle_uuid"], + "starvla.component.text.filename": text_filename, + "starvla.component.mmproj.filename": mmproj_filename, + "starvla.qwen.hidden_size": dimensions["qwen_hidden_dim"], + "starvla.qwen.input_embedding_size": 4 * dimensions["qwen_hidden_dim"], + "starvla.qwen.layer_count": qwen["layer_count"], + "starvla.qwen.vocab_size": qwen["vocab_size"], + "starvla.prompt.cot_template": str(vla.get("CoT_prompt", "")), + "starvla.image.count": len(image_names), + "starvla.image.names": image_names, + "starvla.image.processor_min_pixels": processor_size["shortest_edge"], + "starvla.image.processor_max_pixels": processor_size["longest_edge"], + "starvla.image.patch_size": qwen["processor_patch_size"], + "starvla.image.spatial_merge_size": qwen["processor_merge_size"], + "starvla.image.min_token_count": QWEN3VL_MIN_IMAGE_TOKENS, + "starvla.image.max_token_count": QWEN3VL_MAX_IMAGE_TOKENS, + "starvla.action.dimension": dimensions["action_dim"], + "starvla.action.horizon": horizon, + "starvla.action.continuous_dimensions": list(range(dimensions["action_dim"] - 1)), + "starvla.action.binary_dimensions": [dimensions["action_dim"] - 1], + "starvla.pi_v3.dit_width": dimensions["dit_width"], + "starvla.pi_v3.block_count": dimensions["block_count"], + "starvla.pi_v3.projector_count": dimensions["projector_count"], + "starvla.pi_v3.attention_head_count": diffusion["num_attention_heads"], + "starvla.pi_v3.attention_head_dim": diffusion["attention_head_dim"], + "starvla.pi_v3.feed_forward_dim": dimensions["feed_forward_dim"], + "starvla.pi_v3.mlp_hidden_dimension": dimensions["mlp_hidden_dim"], + "starvla.pi_v3.future_token_count": dimensions["future_token_count"], + "starvla.pi_v3.action_position_count": dimensions["max_sequence_length"], + "starvla.pi_v3.no_state_sequence_length": dimensions["future_token_count"] + horizon, + "starvla.pi_v3.timestep_projection_dim": dimensions["timestep_dim"], + "starvla.pi_v3.num_timestep_buckets": timestep_buckets, + "starvla.pi_v3.num_inference_timesteps": num_steps, + "starvla.pi_v3.ada_norm_epsilon": PI_V3_DIT_NORM_EPS, + "starvla.pi_v3.projector_norm_epsilon": PI_V3_PROJECTOR_NORM_EPS, + "starvla.pi_v3.euler_dt": 1.0 / num_steps, + } + metadata.update( + normalization_metadata( + _load_json(policy_dir / "dataset_statistics.json"), + dimensions["action_dim"], + ) + ) + return metadata + + +def convert_oft_policy( + policy_dir: Path, + hf_dir: Path, + surgery_manifest_path: Path, + output: Path, + catalog_path: Path, + dtype: str, + text_filename: str, + mmproj_filename: str, +) -> None: + catalog = load_catalog(catalog_path) + surgery_manifest = _load_json(surgery_manifest_path) + variant = get_variant(catalog, str(surgery_manifest.get("variant", ""))) + if variant.get("framework") != "oft": + raise StarVLAError( + f"surgery variant {surgery_manifest.get('variant')!r} is not an OFT policy" + ) + validate_official_surgery_manifest(surgery_manifest, variant, catalog) + verify_staged_assets(hf_dir, surgery_manifest.get("qwen_assets", {}), component="Qwen") + verify_staged_assets(policy_dir, surgery_manifest.get("policy_assets", {}), component="policy") + verify_staged_tensors_against_checkpoint( + policy_dir, + surgery_manifest.get("policy_output", {}), + surgery_manifest, + variant, + component="policy", + ) + + tensors = load_policy_tensors(policy_dir) + dimensions = validate_oft_tensors(tensors) + action_token_id = resolve_action_token_id(hf_dir) + metadata = build_oft_metadata( + policy_dir, + hf_dir, + variant, + surgery_manifest, + dimensions, + action_token_id, + text_filename, + mmproj_filename, + ) + + pi0_writer_dir = Path(__file__).resolve().parents[1] / "pi0" + sys.path.insert(0, str(pi0_writer_dir)) + try: + from gguf_writer import write_gguf_arrays + except ImportError as exc: + raise StarVLAError(f"failed to import repository GGUF writer adapter: {exc}") from exc + + def arrays(): + for source_name, destination_name in OFT_TENSOR_MAP.items(): + tensor = tensors[source_name] + array = tensor.detach().float().cpu().numpy() + yield destination_name, [int(dim) for dim in tensor.shape], np.asarray(array), dtype + + _write_gguf_arrays_no_overwrite(output, metadata, arrays(), write_gguf_arrays) + + +def convert_groot_policy( + policy_dir: Path, + hf_dir: Path, + surgery_manifest_path: Path, + output: Path, + catalog_path: Path, + dtype: str, + text_filename: str, + mmproj_filename: str, +) -> None: + catalog = load_catalog(catalog_path) + surgery_manifest = _load_json(surgery_manifest_path) + variant = get_variant(catalog, str(surgery_manifest.get("variant", ""))) + if variant.get("framework") != "groot": + raise StarVLAError( + f"surgery variant {surgery_manifest.get('variant')!r} is not a GR00T policy" + ) + validate_official_surgery_manifest(surgery_manifest, variant, catalog) + verify_staged_assets(hf_dir, surgery_manifest.get("qwen_assets", {}), component="Qwen") + verify_staged_assets(policy_dir, surgery_manifest.get("policy_assets", {}), component="policy") + verify_staged_tensors_against_checkpoint( + policy_dir, + surgery_manifest.get("policy_output", {}), + surgery_manifest, + variant, + component="policy", + ) + + tensors = load_policy_tensors(policy_dir) + dimensions = validate_groot_tensors(tensors) + metadata = build_groot_metadata( + policy_dir, + hf_dir, + variant, + surgery_manifest, + dimensions, + text_filename, + mmproj_filename, + ) + + pi0_writer_dir = Path(__file__).resolve().parents[1] / "pi0" + sys.path.insert(0, str(pi0_writer_dir)) + try: + from gguf_writer import write_gguf_arrays + except ImportError as exc: + raise StarVLAError(f"failed to import repository GGUF writer adapter: {exc}") from exc + + def arrays(): + for source_name, destination_name in GROOT_TENSOR_MAP.items(): + tensor = tensors[source_name] + array = tensor.detach().float().cpu().numpy() + yield destination_name, [int(dim) for dim in tensor.shape], np.asarray(array), dtype + + _write_gguf_arrays_no_overwrite(output, metadata, arrays(), write_gguf_arrays) + + +def convert_pi_policy( + policy_dir: Path, + hf_dir: Path, + surgery_manifest_path: Path, + output: Path, + catalog_path: Path, + dtype: str, + text_filename: str, + mmproj_filename: str, +) -> None: + catalog = load_catalog(catalog_path) + surgery_manifest = _load_json(surgery_manifest_path) + variant = get_variant(catalog, str(surgery_manifest.get("variant", ""))) + if variant.get("framework") != "pi" or variant.get("backbone") != "qwen2_5_vl": + raise StarVLAError( + f"surgery variant {surgery_manifest.get('variant')!r} is not a Qwen2.5 legacy PI policy" + ) + validate_official_surgery_manifest(surgery_manifest, variant, catalog) + verify_staged_assets(hf_dir, surgery_manifest.get("qwen_assets", {}), component="Qwen") + verify_staged_assets(policy_dir, surgery_manifest.get("policy_assets", {}), component="policy") + verify_staged_tensors_against_checkpoint( + policy_dir, + surgery_manifest.get("policy_output", {}), + surgery_manifest, + variant, + component="policy", + ) + + tensors = load_policy_tensors(policy_dir) + dimensions = validate_pi_tensors(tensors) + metadata = build_pi_metadata( + policy_dir, + hf_dir, + variant, + surgery_manifest, + dimensions, + text_filename, + mmproj_filename, + ) + + pi0_writer_dir = Path(__file__).resolve().parents[1] / "pi0" + sys.path.insert(0, str(pi0_writer_dir)) + try: + from gguf_writer import write_gguf_arrays + except ImportError as exc: + raise StarVLAError(f"failed to import repository GGUF writer adapter: {exc}") from exc + + def arrays(): + for source_name, destination_name in PI_TENSOR_MAP.items(): + tensor = tensors[source_name] + array = tensor.detach().float().cpu().numpy() + yield destination_name, [int(dim) for dim in tensor.shape], np.asarray(array), dtype + + _write_gguf_arrays_no_overwrite(output, metadata, arrays(), write_gguf_arrays) + + +def convert_pi_v3_policy( + policy_dir: Path, + hf_dir: Path, + surgery_manifest_path: Path, + output: Path, + catalog_path: Path, + dtype: str, + text_filename: str, + mmproj_filename: str, +) -> None: + catalog = load_catalog(catalog_path) + variant = get_variant(catalog, "pi_v3") + surgery_manifest = _load_json(surgery_manifest_path) + validate_official_surgery_manifest(surgery_manifest, variant, catalog) + verify_staged_assets(hf_dir, surgery_manifest.get("qwen_assets", {}), component="Qwen") + verify_staged_assets(policy_dir, surgery_manifest.get("policy_assets", {}), component="policy") + verify_staged_tensors_against_checkpoint( + policy_dir, + surgery_manifest.get("policy_output", {}), + surgery_manifest, + variant, + component="policy", + ) + + tensors = load_policy_tensors(policy_dir) + dimensions = validate_pi_v3_tensors(tensors) + metadata = build_pi_v3_metadata( + policy_dir, + hf_dir, + variant, + surgery_manifest, + dimensions, + text_filename, + mmproj_filename, + ) + + pi0_writer_dir = Path(__file__).resolve().parents[1] / "pi0" + sys.path.insert(0, str(pi0_writer_dir)) + try: + from gguf_writer import write_gguf_arrays + except ImportError as exc: + raise StarVLAError(f"failed to import repository GGUF writer adapter: {exc}") from exc + + def arrays(): + for source_name, destination_name in PI_V3_TENSOR_MAP.items(): + tensor = tensors[source_name] + array = tensor.detach().float().cpu().numpy() + yield destination_name, [int(dim) for dim in tensor.shape], np.asarray(array), dtype + + _write_gguf_arrays_no_overwrite(output, metadata, arrays(), write_gguf_arrays) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--variant", + default="oft", + choices=( + "oft", + "groot", + "pi_v3", + "qwen25_oft", + "qwen25_groot", + "qwen25_pi", + ), + ) + parser.add_argument("--policy-dir", type=Path, required=True) + parser.add_argument("--hf-dir", type=Path, required=True) + parser.add_argument("--surgery-manifest", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--catalog", type=Path, default=DEFAULT_CATALOG) + parser.add_argument( + "--dtype", + choices=("fp32", "f16", "bf16"), + default=DEFAULT_POLICY_DTYPE, + ) + parser.add_argument("--text-filename") + parser.add_argument("--mmproj-filename") + return parser.parse_args() + + +def main() -> int: + args = parse_args() + try: + if args.output.exists() or args.output.is_symlink(): + raise StarVLAError(f"refusing to overwrite existing output: {args.output}") + text_filename = args.text_filename or default_text_filename( + args.variant, DEFAULT_TEXT_DTYPE + ) + mmproj_filename = args.mmproj_filename or default_mmproj_filename( + args.variant, DEFAULT_MMPROJ_DTYPE + ) + converters = { + "oft": convert_oft_policy, + "groot": convert_groot_policy, + "pi_v3": convert_pi_v3_policy, + "qwen25_oft": convert_oft_policy, + "qwen25_groot": convert_groot_policy, + "qwen25_pi": convert_pi_policy, + } + converter = converters[args.variant] + converter( + policy_dir=args.policy_dir, + hf_dir=args.hf_dir, + surgery_manifest_path=args.surgery_manifest, + output=args.output, + catalog_path=args.catalog, + dtype=args.dtype, + text_filename=text_filename, + mmproj_filename=mmproj_filename, + ) + print(f"policy GGUF: {args.output}") + return 0 + except (StarVLAError, OSError, json.JSONDecodeError, ValueError, TypeError) as exc: + print(f"error: {exc}", file=sys.stderr) + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/hf2gguf/starvla/convert_starvla_qwen_to_gguf.py b/tools/hf2gguf/starvla/convert_starvla_qwen_to_gguf.py new file mode 100755 index 0000000..8a24289 --- /dev/null +++ b/tools/hf2gguf/starvla/convert_starvla_qwen_to_gguf.py @@ -0,0 +1,286 @@ +#!/usr/bin/env python3 +"""Invoke the pinned llama.cpp converter for StarVLA Qwen-VL text and mmproj.""" + +from __future__ import annotations + +import argparse +import json +import re +import subprocess +import sys +from pathlib import Path + +from starvla_checkpoint import ( + DEFAULT_CATALOG, + DEFAULT_MMPROJ_DTYPE, + DEFAULT_TEXT_DTYPE, + StarVLAError, + atomic_write_json, + default_mmproj_filename, + default_text_filename, + get_variant, + load_catalog, + validate_official_surgery_manifest, + verify_staged_assets, + verify_staged_tensors_against_checkpoint, +) + + +REPOSITORY_ROOT = Path(__file__).resolve().parents[3] +LLAMA_ROOT = REPOSITORY_ROOT / "third_party" / "llama.cpp" +LLAMA_CONVERTER = LLAMA_ROOT / "convert_hf_to_gguf.py" +LLAMA_GGUF_PY = LLAMA_ROOT / "gguf-py" +PINNED_REVISION_RE = re.compile(r"[0-9a-f]{40}") + + +def git_revision(path: Path) -> str: + try: + result = subprocess.run( + ["git", "-C", str(path), "rev-parse", "HEAD"], + check=True, + capture_output=True, + text=True, + ) + except (OSError, subprocess.CalledProcessError) as exc: + raise StarVLAError(f"failed to resolve git revision for {path}: {exc}") from exc + return result.stdout.strip() + + +def git_worktree_changes(path: Path) -> str: + try: + result = subprocess.run( + [ + "git", + "-C", + str(path), + "status", + "--porcelain=v1", + "--untracked-files=all", + ], + check=True, + capture_output=True, + text=True, + ) + except (OSError, subprocess.CalledProcessError) as exc: + raise StarVLAError(f"failed to inspect git worktree for {path}: {exc}") from exc + return result.stdout.strip() + + +def canonical_llama_root(path: Path) -> Path: + """Require an explicit, canonical llama.cpp checkout root with converter sources.""" + if not path.is_absolute(): + raise StarVLAError(f"llama.cpp root must be an absolute canonical directory: {path}") + try: + canonical = path.resolve(strict=True) + except (OSError, RuntimeError) as exc: + raise StarVLAError(f"failed to resolve llama.cpp root {path}: {exc}") from exc + if canonical != path or not canonical.is_dir(): + raise StarVLAError(f"llama.cpp root must be an absolute canonical directory: {path}") + + converter = canonical / "convert_hf_to_gguf.py" + gguf_py = canonical / "gguf-py" + if not converter.is_file(): + raise StarVLAError(f"missing llama.cpp converter: {converter}") + if not gguf_py.is_dir(): + raise StarVLAError(f"missing llama.cpp gguf-py directory: {gguf_py}") + return canonical + + +def verify_llama_checkout(path: Path, expected_revision: str) -> Path: + """Verify that path is the clean root of the exact manifest-pinned checkout.""" + root = canonical_llama_root(path) + if PINNED_REVISION_RE.fullmatch(expected_revision) is None: + raise StarVLAError( + f"manifest contains an invalid pinned llama.cpp revision: {expected_revision!r}" + ) + + try: + result = subprocess.run( + ["git", "-C", str(root), "rev-parse", "--show-toplevel"], + check=True, + capture_output=True, + text=True, + ) + except (OSError, subprocess.CalledProcessError) as exc: + raise StarVLAError(f"failed to resolve git root for {root}: {exc}") from exc + try: + git_root = Path(result.stdout.strip()).resolve(strict=True) + except (OSError, RuntimeError) as exc: + raise StarVLAError(f"failed to resolve git root reported for {root}: {exc}") from exc + if git_root != root: + raise StarVLAError( + f"llama.cpp root must be the canonical Git worktree root: expected {root}, got {git_root}" + ) + + actual_revision = git_revision(root) + if actual_revision != expected_revision: + raise StarVLAError( + f"llama.cpp revision mismatch: expected {expected_revision}, got {actual_revision}; " + "update the pinned catalog and regenerate golden data before converting" + ) + worktree_changes = git_worktree_changes(root) + if worktree_changes: + raise StarVLAError( + "llama.cpp has tracked or untracked worktree changes; " + f"use the clean pinned revision for official conversion:\n{worktree_changes}" + ) + return root + + +def build_commands( + python: str, + hf_dir: Path, + text_output: Path, + mmproj_output: Path, + text_metadata: Path, + mmproj_metadata: Path, + text_dtype: str, + mmproj_dtype: str, + *, + llama_root: Path = LLAMA_ROOT, +) -> list[list[str]]: + # Isolated mode excludes the working directory, PYTHONPATH and user site + # from imports while the pinned converter adds its own gguf-py directory. + converter = llama_root / "convert_hf_to_gguf.py" + common = [python, "-I", str(converter), str(hf_dir)] + return [ + common + + [ + "--outfile", + str(text_output), + "--outtype", + text_dtype, + "--metadata", + str(text_metadata), + ], + common + + [ + "--outfile", + str(mmproj_output), + "--outtype", + mmproj_dtype, + "--metadata", + str(mmproj_metadata), + "--mmproj", + ], + ] + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--hf-dir", type=Path, required=True) + parser.add_argument("--surgery-manifest", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--catalog", type=Path, default=DEFAULT_CATALOG) + parser.add_argument("--llama-root", type=Path, default=LLAMA_ROOT) + parser.add_argument("--text-filename") + parser.add_argument("--mmproj-filename") + parser.add_argument( + "--text-dtype", + choices=("f32", "f16", "bf16", "q8_0"), + default=DEFAULT_TEXT_DTYPE, + ) + parser.add_argument( + "--mmproj-dtype", + choices=("f32", "f16", "bf16", "q8_0"), + default=DEFAULT_MMPROJ_DTYPE, + ) + parser.add_argument("--python", default=sys.executable) + parser.add_argument("--dry-run", action="store_true") + return parser.parse_args() + + +def main() -> int: + args = parse_args() + try: + if not args.hf_dir.is_dir(): + raise StarVLAError(f"missing HF staging directory: {args.hf_dir}") + try: + manifest = json.loads(args.surgery_manifest.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise StarVLAError(f"failed to load surgery manifest {args.surgery_manifest}: {exc}") from exc + + catalog = load_catalog(args.catalog) + variant_name = str(manifest.get("variant", "")) + variant = get_variant(catalog, variant_name) + validate_official_surgery_manifest(manifest, variant, catalog) + verify_staged_assets(args.hf_dir, manifest.get("qwen_assets", {}), component="Qwen") + verify_staged_tensors_against_checkpoint( + args.hf_dir, + manifest.get("vlm_output", {}), + manifest, + variant, + component="vlm", + ) + + expected_revision = str(manifest.get("source", {}).get("llama_cpp_revision", "")) + llama_root = verify_llama_checkout(args.llama_root, expected_revision) + + args.output_dir.mkdir(parents=True, exist_ok=True) + text_filename = args.text_filename or default_text_filename(variant_name, args.text_dtype) + mmproj_filename = args.mmproj_filename or default_mmproj_filename( + variant_name, args.mmproj_dtype + ) + text_output = args.output_dir / text_filename + mmproj_output = args.output_dir / mmproj_filename + text_metadata = args.output_dir / "text-metadata.json" + mmproj_metadata = args.output_dir / "mmproj-metadata.json" + bundle_uuid = str(manifest["bundle_uuid"]) + source = manifest["source"] + backbone = str(manifest.get("backbone", variant.get("backbone", "qwen3_vl"))) + backbone_label = { + "qwen3_vl": "Qwen3-VL", + "qwen2_5_vl": "Qwen2.5-VL", + }.get(backbone) + if backbone_label is None: + raise StarVLAError(f"unsupported StarVLA Qwen backbone: {backbone!r}") + common_metadata = { + "general.source.uuid": bundle_uuid, + "general.source.url": f"https://huggingface.co/{source['repo_id']}/tree/{source['revision']}", + "general.finetune": f"starvla-{manifest['variant']}", + } + atomic_write_json( + text_metadata, + { + **common_metadata, + "general.name": f"StarVLA {backbone_label} {manifest['variant']} text", + }, + ) + atomic_write_json( + mmproj_metadata, + { + **common_metadata, + "general.name": f"StarVLA {backbone_label} {manifest['variant']} mmproj", + }, + ) + + commands = build_commands( + args.python, + args.hf_dir, + text_output, + mmproj_output, + text_metadata, + mmproj_metadata, + args.text_dtype, + args.mmproj_dtype, + llama_root=llama_root, + ) + if args.dry_run: + print(json.dumps(commands, indent=2)) + return 0 + + for command in commands: + subprocess.run(command, check=True, cwd=REPOSITORY_ROOT) + for output in (text_output, mmproj_output): + if not output.is_file() or output.stat().st_size == 0: + raise StarVLAError(f"llama.cpp converter did not create the expected output: {output}") + print(f"text GGUF: {text_output}") + print(f"mmproj GGUF: {mmproj_output}") + return 0 + except (StarVLAError, OSError, subprocess.CalledProcessError, KeyError) as exc: + print(f"error: {exc}", file=sys.stderr) + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/hf2gguf/starvla/download_starvla.py b/tools/hf2gguf/starvla/download_starvla.py new file mode 100755 index 0000000..ceb712f --- /dev/null +++ b/tools/hf2gguf/starvla/download_starvla.py @@ -0,0 +1,422 @@ +#!/usr/bin/env python3 +"""Download pinned StarVLA sources and shared tokenizer assets.""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path +from typing import Any, Sequence + +from starvla_checkpoint import ( + DEFAULT_CATALOG, + DEFAULT_QWEN_ASSET, + SUPPORTED_BACKBONES, + StarVLAError, + atomic_write_json, + get_variant, + load_catalog, + sha256_file, +) + + +DEFAULT_BACKBONE = "qwen3_vl" +DEFAULT_TARGET_MATRIX = Path(__file__).with_name("release_targets.json") + + +def load_target_matrix(path: Path | str = DEFAULT_TARGET_MATRIX) -> dict[str, Any]: + import json + + matrix_path = Path(path) + try: + matrix = json.loads(matrix_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise StarVLAError(f"failed to load release target matrix {matrix_path}: {exc}") from exc + if not isinstance(matrix, dict) or matrix.get("schema_version") != 1: + raise StarVLAError("unsupported StarVLA release target matrix") + return matrix + + +def validate_target_matrix( + catalog: dict[str, Any], matrix: dict[str, Any] +) -> None: + variants = catalog["variants"] + expected_backbones = set(SUPPORTED_BACKBONES) + seen: set[str] = set() + for tier in ("targets", "experimental"): + groups = matrix.get(tier) + if not isinstance(groups, dict) or set(groups) != expected_backbones: + raise StarVLAError( + f"release target matrix {tier} must cover every supported backbone exactly" + ) + for backbone, names in groups.items(): + if ( + not isinstance(names, list) + or any(not isinstance(name, str) for name in names) + or len(names) != len(set(names)) + ): + raise StarVLAError( + f"release target matrix {tier}.{backbone} must be a unique list" + ) + if tier == "targets" and not names: + raise StarVLAError( + f"release target matrix targets.{backbone} cannot be empty" + ) + for name in names: + entry = variants.get(name) + if not isinstance(entry, dict): + raise StarVLAError(f"release target {name!r} is not a catalog variant") + if variant_backbone(entry) != backbone: + raise StarVLAError(f"release target {name!r} has the wrong backbone") + if name in seen: + raise StarVLAError(f"release target {name!r} occurs more than once") + if ( + entry.get("status") != "official_policy" + or entry.get("checkpoint") is None + ): + raise StarVLAError( + f"release target {name!r} is not an official policy checkpoint" + ) + seen.add(name) + policy_variants = { + name + for name, entry in variants.items() + if entry.get("status") == "official_policy" and entry.get("checkpoint") is not None + } + if seen != policy_variants: + raise StarVLAError( + "release target matrix must classify every policy variant exactly once" + ) + + +def destination_for(root: Path, entry: dict[str, Any]) -> Path: + return root / str(entry["directory"]) / str(entry["revision"]) + + +def variant_backbone(entry: dict[str, Any]) -> str: + return str(entry.get("backbone", DEFAULT_BACKBONE)) + + +def available_backbones(catalog: dict[str, Any]) -> list[str]: + return list( + dict.fromkeys( + variant_backbone(entry) + for entry in catalog["variants"].values() + ) + ) + + +def resolve_variant_keys( + catalog: dict[str, Any], + requested: Sequence[str] | None, + requested_backbone: str | None, + target_matrix: dict[str, Any] | None = None, +) -> tuple[str, list[str]]: + variants = catalog["variants"] + backbones = available_backbones(catalog) + if requested_backbone is not None: + backbone = requested_backbone + else: + direct_backbones = { + variant_backbone(variants[name]) + for name in (requested or ()) + if name in variants + } + if len(direct_backbones) > 1: + raise StarVLAError( + "requested variants span multiple backbones; select one with --backbone" + ) + backbone = next(iter(direct_backbones), DEFAULT_BACKBONE) + if backbone not in backbones: + raise StarVLAError( + f"unknown StarVLA backbone {backbone!r}; expected one of {backbones}" + ) + + candidates = { + name: entry + for name, entry in variants.items() + if variant_backbone(entry) == backbone + } + tokens = list(requested or ("oft",)) + if "all" in tokens or "catalog-all" in tokens: + if len(tokens) != 1: + raise StarVLAError( + "--variant all/catalog-all cannot be combined with another variant" + ) + if tokens[0] == "all": + matrix = target_matrix or load_target_matrix() + validate_target_matrix(catalog, matrix) + return backbone, list(matrix["targets"][backbone]) + return backbone, list(candidates) + + selected: list[str] = [] + for token in tokens: + if token in candidates: + key = token + else: + matches = [ + name + for name, entry in candidates.items() + if entry.get("framework") == token + ] + if not matches: + accepted = sorted( + { + *candidates, + *(str(entry["framework"]) for entry in candidates.values()), + "all", + "catalog-all", + } + ) + raise StarVLAError( + f"variant {token!r} is not available for backbone {backbone!r}; " + f"expected one of {accepted}" + ) + if len(matches) != 1: + raise StarVLAError( + f"framework alias {token!r} is ambiguous for backbone {backbone!r}; " + f"use one of {matches}" + ) + key = matches[0] + if key not in selected: + selected.append(key) + return backbone, selected + + +def required_shared_assets( + catalog: dict[str, Any], + variant_keys: Sequence[str], +) -> list[str]: + names: list[str] = [] + has_fast = False + for variant_key in variant_keys: + entry = get_variant(catalog, variant_key) + qwen_asset = str(entry.get("qwen_asset", DEFAULT_QWEN_ASSET)) + if qwen_asset not in catalog["shared_assets"]: + raise StarVLAError( + f"variant {variant_key!r} references unknown Qwen asset {qwen_asset!r}" + ) + if qwen_asset not in names: + names.append(qwen_asset) + has_fast = has_fast or entry.get("framework") == "fast" + if has_fast and "fast_codec" not in names: + names.append("fast_codec") + return names + + +def download_entry( + entry: dict[str, Any], + root: Path, + files: list[str], + *, + dry_run: bool, + local_files_only: bool, + force_download: bool, +) -> dict[str, Any]: + destination = destination_for(root, entry) + result = { + "repo_id": entry["repo_id"], + "revision": entry["revision"], + "directory": str(destination), + "requested_files": files, + "files": [], + } + if dry_run: + return result + + try: + from huggingface_hub import snapshot_download + except ImportError as exc: + raise StarVLAError("huggingface_hub is required; install tools/hf2gguf/environment.yaml") from exc + + destination.mkdir(parents=True, exist_ok=True) + try: + snapshot_download( + repo_id=str(entry["repo_id"]), + revision=str(entry["revision"]), + allow_patterns=files, + local_dir=destination, + local_files_only=local_files_only, + force_download=force_download, + ) + except Exception as exc: + raise StarVLAError(f"failed to download {entry['repo_id']}@{entry['revision']}: {exc}") from exc + + missing = [relative for relative in files if not (destination / relative).is_file()] + if missing: + raise StarVLAError(f"download completed with missing files in {destination}: {missing}") + + expected_records = dict(entry.get("file_hashes", {})) + expected_records.update(entry.get("optional_weight_hashes", {})) + checkpoint = entry.get("checkpoint") + if checkpoint is not None: + expected_records[str(checkpoint["path"])] = checkpoint + for relative in sorted(files): + path = destination / relative + expected = expected_records.get(relative) + if expected is None: + raise StarVLAError(f"catalog has no pinned size/SHA256 for requested file: {relative}") + record = {"path": relative, "size": path.stat().st_size, "sha256": sha256_file(path)} + if record["size"] != expected["size"] or record["sha256"] != expected["sha256"]: + raise StarVLAError( + f"downloaded file size/SHA256 mismatch for {path}: " + f"expected {expected['size']}/{expected['sha256']}, " + f"got {record['size']}/{record['sha256']}" + ) + result["files"].append(record) + return result + + +def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--variant", + action="append", + help=( + "catalog variant key or framework alias to download " + "(repeatable; 'all' selects the release matrix, 'catalog-all' also " + "includes experimental entries; default: OFT for the selected backbone)" + ), + ) + parser.add_argument( + "--backbone", + help=( + "backbone selector from the catalog " + f"(default: infer from exact variant keys, otherwise {DEFAULT_BACKBONE})" + ), + ) + parser.add_argument("--root", type=Path, default=Path("ckpts/starvla/sources")) + parser.add_argument("--catalog", type=Path, default=DEFAULT_CATALOG) + parser.add_argument( + "--target-matrix", type=Path, default=DEFAULT_TARGET_MATRIX, + help="release/experimental support matrix used by --variant all", + ) + parser.add_argument( + "--metadata-only", + action="store_true", + help="skip policy checkpoints and all optional base weights", + ) + parser.add_argument( + "--include-base-weights", + action="store_true", + help="download optional safetensor shards for the selected Qwen base assets", + ) + parser.add_argument( + "--include-fast-weights", + action="store_true", + help=( + "download the action-ready base weights for a selected FAST variant " + "(kept for compatibility; the policy checkpoint is downloaded separately)" + ), + ) + parser.add_argument("--no-shared-assets", action="store_true") + parser.add_argument("--local-files-only", action="store_true") + parser.add_argument("--force-download", action="store_true") + parser.add_argument("--dry-run", action="store_true") + return parser.parse_args(argv) + + +def main() -> int: + args = parse_args() + try: + catalog = load_catalog(args.catalog) + target_matrix = load_target_matrix(args.target_matrix) + validate_target_matrix(catalog, target_matrix) + backbone, variants = resolve_variant_keys( + catalog, args.variant, args.backbone, target_matrix + ) + + manifest: dict[str, Any] = { + "schema_version": 1, + "catalog": str(args.catalog.resolve()), + "target_matrix": str(args.target_matrix.resolve()), + "target_matrix_sha256": sha256_file(args.target_matrix), + "source_revisions": catalog["source_revisions"], + "backbone": backbone, + "variants": variants, + "metadata_only": bool(args.metadata_only), + "downloads": {}, + } + + if not args.no_shared_assets: + shared_names = required_shared_assets(catalog, variants) + variant_entries = { + variant: get_variant(catalog, variant) + for variant in variants + } + fast_qwen_assets = { + str(entry.get("qwen_asset", DEFAULT_QWEN_ASSET)) + for entry in variant_entries.values() + if entry.get("framework") == "fast" + } + for name in shared_names: + raw_entry = catalog["shared_assets"][name] + entry = dict(raw_entry) + files = list(entry["files"]) + include_optional_weights = ( + args.include_base_weights + or (args.include_fast_weights and name in fast_qwen_assets) + ) + if include_optional_weights and not args.metadata_only: + files.extend(entry.get("optional_weight_files", [])) + manifest["downloads"][f"asset:{name}"] = download_entry( + entry, + args.root, + files, + dry_run=args.dry_run, + local_files_only=args.local_files_only, + force_download=args.force_download, + ) + + for variant in variants: + entry = get_variant(catalog, variant) + files = list(entry.get("files", [])) + checkpoint = entry.get("checkpoint") + if checkpoint is not None and not args.metadata_only: + files.append(str(checkpoint["path"])) + include_variant_weights = ( + args.include_base_weights + or (entry.get("framework") == "fast" and args.include_fast_weights) + ) + if include_variant_weights and not args.metadata_only: + files.extend(entry.get("optional_weight_files", [])) + download = download_entry( + entry, + args.root, + files, + dry_run=args.dry_run, + local_files_only=args.local_files_only, + force_download=args.force_download, + ) + manifest["downloads"][f"variant:{variant}"] = download + + if checkpoint is not None and not args.metadata_only and not args.dry_run: + record = next( + (item for item in download["files"] if item["path"] == checkpoint["path"]), + None, + ) + if record is None: + raise StarVLAError(f"download manifest has no checkpoint record for {checkpoint['path']}") + if record["size"] != checkpoint["size"] or record["sha256"] != checkpoint["sha256"]: + raise StarVLAError( + f"checkpoint verification failed for {checkpoint['path']}: " + f"expected size/hash {checkpoint['size']}/{checkpoint['sha256']}, " + f"got {record['size']}/{record['sha256']}" + ) + + if args.dry_run: + import json + + print(json.dumps(manifest, indent=2, sort_keys=True)) + else: + manifest_path = args.root / "download_manifest.json" + atomic_write_json(manifest_path, manifest) + print(f"download manifest: {manifest_path}") + return 0 + except StarVLAError as exc: + print(f"error: {exc}", file=sys.stderr) + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/hf2gguf/starvla/environment.yaml b/tools/hf2gguf/starvla/environment.yaml new file mode 100644 index 0000000..2deb0dd --- /dev/null +++ b/tools/hf2gguf/starvla/environment.yaml @@ -0,0 +1,14 @@ +name: starvla_gguf_converter +channels: + - conda-forge +dependencies: + - python=3.11 + - pip + - pip: + - torch + - numpy + - safetensors + - sentencepiece + - transformers==4.57.0 + - huggingface_hub>=0.36.0 + - pyyaml diff --git a/tools/hf2gguf/starvla/inspect_starvla_checkpoint.py b/tools/hf2gguf/starvla/inspect_starvla_checkpoint.py new file mode 100755 index 0000000..822ebf0 --- /dev/null +++ b/tools/hf2gguf/starvla/inspect_starvla_checkpoint.py @@ -0,0 +1,177 @@ +#!/usr/bin/env python3 +"""Inspect a StarVLA .pt checkpoint without materializing copied weights.""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path +from typing import Any + +from starvla_checkpoint import ( + DEFAULT_CATALOG, + StarVLAError, + atomic_write_json, + build_inventory, + get_variant, + inventory_summary, + load_catalog, + load_checkpoint_state, + resolve_effective_config, + sha256_file, + verify_catalog_files, +) + + +def _load_structured(path: Path) -> dict[str, Any]: + if path.suffix == ".json": + return json.loads(path.read_text(encoding="utf-8")) + try: + import yaml + except ImportError as exc: + raise StarVLAError("PyYAML is required to inspect StarVLA YAML configs") from exc + value = yaml.safe_load(path.read_text(encoding="utf-8")) + if not isinstance(value, dict): + raise StarVLAError(f"expected an object in config file {path}") + return value + + +def _flatten(value: Any, prefix: str = "") -> dict[str, Any]: + if not isinstance(value, dict): + return {prefix: value} + flattened: dict[str, Any] = {} + for key, child in value.items(): + path = f"{prefix}.{key}" if prefix else str(key) + flattened.update(_flatten(child, path)) + return flattened + + +def inspect_config_candidates(source_dir: Path | None) -> dict[str, Any]: + if source_dir is None: + return {"files": {}, "conflicts": {}} + configs: dict[str, dict[str, Any]] = {} + for name in ("config.json", "config.yaml", "config.full.yaml"): + path = source_dir / name + if path.is_file(): + configs[name] = _load_structured(path) + + flattened = {name: _flatten(value) for name, value in configs.items()} + all_keys = sorted({key for values in flattened.values() for key in values}) + conflicts: dict[str, dict[str, Any]] = {} + for key in all_keys: + observed = {name: values[key] for name, values in flattened.items() if key in values} + serialized = {json.dumps(value, sort_keys=True) for value in observed.values()} + if len(observed) > 1 and len(serialized) > 1: + conflicts[key] = observed + return {"files": configs, "conflicts": conflicts} + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("checkpoint", type=Path) + parser.add_argument( + "--variant", + required=True, + choices=( + "oft", + "groot", + "pi_v3", + "qwen25_oft", + "qwen25_groot", + "qwen25_pi", + ), + ) + parser.add_argument("--catalog", type=Path, default=DEFAULT_CATALOG) + parser.add_argument("--source-dir", type=Path, help="directory containing config/statistics files") + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--effective-config-output", type=Path) + parser.add_argument("--skip-hash-check", action="store_true") + parser.add_argument( + "--allow-nonofficial-inventory", + action="store_true", + help="do not enforce pinned tensor counts; intended only for synthetic tests", + ) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + try: + effective_config_output = args.effective_config_output or args.output.with_name( + "effective_config.json" + ) + output_targets = [args.output] + if args.source_dir is not None: + output_targets.append(effective_config_output) + if len(set(output_targets)) != len(output_targets): + raise StarVLAError("inspection and effective-config outputs must be different files") + for output_target in output_targets: + if output_target.exists() or output_target.is_symlink(): + raise StarVLAError(f"refusing to overwrite existing output: {output_target}") + + catalog = load_catalog(args.catalog) + variant = get_variant(catalog, args.variant) + if not args.checkpoint.is_file(): + raise StarVLAError(f"checkpoint does not exist: {args.checkpoint}") + actual_checkpoint = { + "path": str(args.checkpoint.resolve()), + "size": args.checkpoint.stat().st_size, + "sha256": sha256_file(args.checkpoint), + } + expected_checkpoint = variant["checkpoint"] + checkpoint_verified = ( + actual_checkpoint["size"] == expected_checkpoint["size"] + and actual_checkpoint["sha256"] == expected_checkpoint["sha256"] + ) + if not args.skip_hash_check and not checkpoint_verified: + raise StarVLAError( + f"checkpoint size/SHA256 mismatch for {args.checkpoint}: " + f"expected {expected_checkpoint['size']}/{expected_checkpoint['sha256']}, " + f"got {actual_checkpoint['size']}/{actual_checkpoint['sha256']}" + ) + source_assets_verified = False + if args.source_dir is not None: + verify_catalog_files(args.source_dir, variant) + source_assets_verified = True + state_dict = load_checkpoint_state(args.checkpoint) + records = build_inventory( + state_dict, + variant, + enforce_expected=not args.allow_nonofficial_inventory, + ) + effective_config = ( + resolve_effective_config(args.source_dir, args.variant, variant) + if args.source_dir + else None + ) + if effective_config is not None: + atomic_write_json(effective_config_output, effective_config, overwrite=False) + result = { + "schema_version": 1, + "variant": args.variant, + "model_type": variant["model_type"], + "source": { + "repo_id": variant["repo_id"], + "revision": variant["revision"], + "catalog_checkpoint": expected_checkpoint, + "input_checkpoint": actual_checkpoint, + "checkpoint_verification": "verified" if checkpoint_verified else "skipped_nonofficial", + "source_assets_verified": source_assets_verified, + }, + "summary": inventory_summary(records), + "config_candidates": inspect_config_candidates(args.source_dir), + "effective_config": str(effective_config_output) if effective_config is not None else None, + "tensors": [record.to_json() for record in records], + } + atomic_write_json(args.output, result, overwrite=False) + print(json.dumps(result["summary"], indent=2, sort_keys=True)) + print(f"inventory: {args.output}") + return 0 + except (StarVLAError, OSError, json.JSONDecodeError) as exc: + print(f"error: {exc}", file=sys.stderr) + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/hf2gguf/starvla/starvla_checkpoint.py b/tools/hf2gguf/starvla/starvla_checkpoint.py new file mode 100755 index 0000000..ef50d79 --- /dev/null +++ b/tools/hf2gguf/starvla/starvla_checkpoint.py @@ -0,0 +1,1152 @@ +#!/usr/bin/env python3 +"""Shared catalog and strict checkpoint inventory helpers for StarVLA.""" + +from __future__ import annotations + +import hashlib +import json +import os +import uuid +from copy import deepcopy +from collections import Counter, defaultdict +from dataclasses import asdict, dataclass +from pathlib import Path, PureWindowsPath +from typing import Any, Mapping, Sequence + + +DEFAULT_CATALOG = Path(__file__).with_name("checkpoint_catalog.json") +DEFAULT_TEXT_DTYPE = "bf16" +DEFAULT_MMPROJ_DTYPE = "bf16" +DEFAULT_POLICY_DTYPE = "fp32" + +STARVLA_ARTIFACT_STEMS = { + "oft": "oft", + "groot": "groot", + "pi_v3": "pi-v3", + "qwen25_oft": "qwen25-oft", + "qwen25_groot": "qwen25-groot", + "qwen25_pi": "qwen25-pi", + "qwen25_fast": "qwen25-fast", +} + +DEFAULT_QWEN_ASSET = "qwen3_vl_4b_instruct" +SUPPORTED_BACKBONES = {"qwen3_vl", "qwen2_5_vl"} +SUPPORTED_FRAMEWORKS = {"oft", "groot", "pi", "pi_v3", "fast"} +GENERATED_QWEN_ASSET_PATHS = {"model.safetensors.index.json"} + +VLM_SOURCE_RULES = ( + ("qwen_vl_interface.model.model.visual.", "visual"), + ("qwen_vl_interface.model.model.language_model.", "text"), + ("qwen_vl_interface.model.lm_head.", "lm_head"), +) + +VLM_DESTINATION_PREFIXES = { + "qwen3_vl": { + "visual": "model.visual.", + "text": "model.language_model.", + "lm_head": "lm_head.", + }, + "qwen2_5_vl": { + "visual": "visual.", + "text": "model.", + "lm_head": "lm_head.", + }, +} + + +class StarVLAError(RuntimeError): + """Raised when a catalog or checkpoint violates the conversion contract.""" + + +def artifact_stem(variant: str) -> str: + try: + return STARVLA_ARTIFACT_STEMS[variant] + except KeyError as exc: + raise StarVLAError(f"unsupported StarVLA artifact variant: {variant!r}") from exc + + +def default_text_filename(variant: str, dtype: str = DEFAULT_TEXT_DTYPE) -> str: + return f"qwen-{artifact_stem(variant)}-{dtype}.gguf" + + +def default_mmproj_filename(variant: str, dtype: str = DEFAULT_MMPROJ_DTYPE) -> str: + return f"mmproj-{artifact_stem(variant)}-{dtype}.gguf" + + +@dataclass(frozen=True) +class TensorRecord: + source_name: str + destination_name: str + component: str + role: str + shape: list[int] + dtype: str + numel: int + nbytes: int + storage_offset: int + storage_alias: str | None = None + + def to_json(self) -> dict[str, Any]: + return asdict(self) + + +def _safe_relative_path(value: Any, *, field: str) -> Path: + if not isinstance(value, str): + raise StarVLAError(f"unsafe relative path in {field}: expected a string, got {value!r}") + if not value or "\x00" in value or "\\" in value: + raise StarVLAError(f"unsafe relative path in {field}: {value!r}") + + path = Path(value) + windows_path = PureWindowsPath(value) + if ( + path.is_absolute() + or windows_path.is_absolute() + or bool(windows_path.drive) + or not path.parts + or any(part in ("", ".", "..") for part in path.parts) + or path.as_posix() != value + ): + raise StarVLAError(f"unsafe relative path in {field}: {value!r}") + return path + + +def _validate_revision(value: Any, *, field: str) -> str: + if ( + not isinstance(value, str) + or len(value) != 40 + or any(ch not in "0123456789abcdef" for ch in value) + ): + raise StarVLAError(f"invalid pinned revision in {field}: expected 40 lowercase hex characters") + return value + + +def _validate_sha256(value: Any, *, field: str) -> str: + if ( + not isinstance(value, str) + or len(value) != 64 + or any(ch not in "0123456789abcdef" for ch in value) + ): + raise StarVLAError(f"invalid SHA256 in {field}") + return value + + +def _validate_positive_size(value: Any, *, field: str) -> int: + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + raise StarVLAError(f"invalid size in {field}: expected a positive integer") + return value + + +def load_catalog(path: Path | str = DEFAULT_CATALOG) -> dict[str, Any]: + catalog_path = Path(path) + try: + catalog = json.loads(catalog_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise StarVLAError(f"failed to load checkpoint catalog {catalog_path}: {exc}") from exc + + if not isinstance(catalog, dict): + raise StarVLAError("checkpoint catalog root must be an object") + if catalog.get("schema_version") != 1: + raise StarVLAError(f"unsupported checkpoint catalog schema: {catalog.get('schema_version')!r}") + + source_revisions = catalog.get("source_revisions") + if not isinstance(source_revisions, dict): + raise StarVLAError("checkpoint catalog source_revisions must be an object") + for source in ("starvla", "llama_cpp"): + _validate_revision(source_revisions.get(source), field=f"source_revisions.{source}") + for source, revision in source_revisions.items(): + _validate_revision(revision, field=f"source_revisions.{source}") + + shared_assets = catalog.get("shared_assets") + if not isinstance(shared_assets, dict): + raise StarVLAError("checkpoint catalog shared_assets must be an object") + variants = catalog.get("variants") + if not isinstance(variants, dict) or not variants: + raise StarVLAError("checkpoint catalog has no variants") + for name, entry in variants.items(): + if not isinstance(entry, dict): + raise StarVLAError(f"catalog variant {name!r} must be an object") + framework = entry.get("framework") + if framework not in SUPPORTED_FRAMEWORKS: + raise StarVLAError( + f"catalog variant {name!r} has unsupported framework={framework!r}" + ) + backbone = entry.get("backbone", "qwen3_vl") + if backbone not in SUPPORTED_BACKBONES: + raise StarVLAError( + f"catalog variant {name!r} has unsupported backbone={backbone!r}" + ) + qwen_asset = entry.get("qwen_asset", DEFAULT_QWEN_ASSET) + if not isinstance(qwen_asset, str) or qwen_asset not in shared_assets: + raise StarVLAError( + f"catalog variant {name!r} references unknown qwen_asset={qwen_asset!r}" + ) + if not isinstance(entry.get("repo_id"), str) or not entry["repo_id"]: + raise StarVLAError(f"catalog variant {name!r} is missing repo_id/revision") + _validate_revision(entry.get("revision"), field=f"variant {name}.revision") + checkpoint = entry.get("checkpoint") + if checkpoint is not None: + if not isinstance(checkpoint, dict): + raise StarVLAError(f"catalog variant {name!r} checkpoint must be an object or null") + _safe_relative_path(checkpoint.get("path"), field=f"variant {name}.checkpoint.path") + _validate_positive_size(checkpoint.get("size"), field=f"variant {name}.checkpoint.size") + _validate_sha256(checkpoint.get("sha256"), field=f"variant {name}.checkpoint.sha256") + entries = { + **{f"shared asset {name}": entry for name, entry in shared_assets.items()}, + **{f"variant {name}": entry for name, entry in variants.items()}, + } + for label, entry in entries.items(): + if not isinstance(entry, dict): + raise StarVLAError(f"catalog {label} must be an object") + _safe_relative_path(entry.get("directory"), field=f"{label}.directory") + if not isinstance(entry.get("repo_id"), str) or not entry["repo_id"]: + raise StarVLAError(f"catalog {label} has invalid repo_id") + _validate_revision(entry.get("revision"), field=f"{label}.revision") + + files = entry.get("files", []) + file_hashes = entry.get("file_hashes", {}) + if not isinstance(files, list) or any(not isinstance(relative, str) for relative in files): + raise StarVLAError(f"catalog {label} has invalid or duplicate files") + for index, relative in enumerate(files): + _safe_relative_path(relative, field=f"{label}.files[{index}]") + if len(files) != len(set(files)): + raise StarVLAError(f"catalog {label} has invalid or duplicate files") + if not isinstance(file_hashes, dict): + raise StarVLAError(f"catalog {label} file_hashes must be an object") + if set(file_hashes) != set(files): + raise StarVLAError(f"catalog {label} file_hashes must cover files exactly") + for relative, record in file_hashes.items(): + if not isinstance(record, dict): + raise StarVLAError(f"catalog {label} has invalid file record for {relative!r}") + _validate_positive_size(record.get("size"), field=f"{label}.file_hashes[{relative!r}].size") + _validate_sha256(record.get("sha256"), field=f"{label}.file_hashes[{relative!r}].sha256") + staged_overrides = entry.get("staged_overrides", {}) + if not isinstance(staged_overrides, dict) or not set(staged_overrides).issubset(files): + raise StarVLAError(f"catalog {label} has invalid staged_overrides") + for relative, record in staged_overrides.items(): + if not isinstance(record, dict): + raise StarVLAError(f"catalog {label} has invalid staged override for {relative!r}") + _validate_positive_size(record.get("size"), field=f"{label}.staged_overrides[{relative!r}].size") + _validate_sha256(record.get("sha256"), field=f"{label}.staged_overrides[{relative!r}].sha256") + optional_files = entry.get("optional_weight_files", []) + optional_hashes = entry.get("optional_weight_hashes", {}) + if not isinstance(optional_files, list) or any(not isinstance(relative, str) for relative in optional_files): + raise StarVLAError(f"catalog {label} has invalid optional_weight_files") + for index, relative in enumerate(optional_files): + _safe_relative_path(relative, field=f"{label}.optional_weight_files[{index}]") + if len(optional_files) != len(set(optional_files)): + raise StarVLAError(f"catalog {label} has duplicate optional weights") + if not isinstance(optional_hashes, dict) or set(optional_hashes) != set(optional_files): + raise StarVLAError(f"catalog {label} optional_weight_hashes must cover optional weights exactly") + for relative, record in optional_hashes.items(): + if not isinstance(record, dict): + raise StarVLAError(f"catalog {label} has invalid optional weight record for {relative!r}") + _validate_positive_size(record.get("size"), field=f"{label}.optional_weight_hashes[{relative!r}].size") + _validate_sha256(record.get("sha256"), field=f"{label}.optional_weight_hashes[{relative!r}].sha256") + return catalog + + +def get_variant(catalog: Mapping[str, Any], variant: str) -> dict[str, Any]: + variants = catalog.get("variants", {}) + if variant not in variants: + raise StarVLAError(f"unknown StarVLA variant {variant!r}; expected one of {sorted(variants)}") + entry = dict(variants[variant]) + entry["_catalog_key"] = variant + return entry + + +def get_qwen_asset( + catalog: Mapping[str, Any], variant_entry: Mapping[str, Any] +) -> tuple[str, dict[str, Any]]: + asset_name = str(variant_entry.get("qwen_asset", DEFAULT_QWEN_ASSET)) + shared_assets = catalog.get("shared_assets", {}) + if asset_name not in shared_assets: + raise StarVLAError( + f"variant {variant_entry.get('_catalog_key', variant_entry.get('framework'))!r} " + f"references unknown Qwen asset {asset_name!r}" + ) + return asset_name, dict(shared_assets[asset_name]) + + +def staged_qwen_asset_hashes(qwen_entry: Mapping[str, Any]) -> dict[str, str]: + """Return immutable assets that survive checkpoint surgery unchanged.""" + return { + relative: qwen_entry.get("staged_overrides", {}).get(relative, record)[ + "sha256" + ] + for relative, record in qwen_entry["file_hashes"].items() + if relative not in GENERATED_QWEN_ASSET_PATHS + } + + +def sha256_file(path: Path, chunk_size: int = 8 * 1024 * 1024) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(chunk_size): + digest.update(chunk) + return digest.hexdigest() + + +def verify_checkpoint_file(path: Path, variant_entry: Mapping[str, Any]) -> None: + checkpoint = variant_entry.get("checkpoint") + if checkpoint is None: + raise StarVLAError(f"variant {variant_entry.get('framework')!r} has no official policy checkpoint") + if not path.is_file(): + raise StarVLAError(f"checkpoint does not exist: {path}") + expected_size = int(checkpoint["size"]) + actual_size = path.stat().st_size + if actual_size != expected_size: + raise StarVLAError(f"checkpoint size mismatch for {path}: expected {expected_size}, got {actual_size}") + expected_hash = str(checkpoint["sha256"]) + actual_hash = sha256_file(path) + if actual_hash != expected_hash: + raise StarVLAError(f"checkpoint SHA256 mismatch for {path}: expected {expected_hash}, got {actual_hash}") + + +def verify_catalog_files(root: Path, entry: Mapping[str, Any]) -> dict[str, str]: + verified = {} + for relative in entry.get("files", []): + path = root / relative + expected = entry["file_hashes"][relative] + if not path.is_file(): + raise StarVLAError(f"missing pinned source asset: {path}") + actual_size = path.stat().st_size + actual_hash = sha256_file(path) + if actual_size != expected["size"] or actual_hash != expected["sha256"]: + raise StarVLAError( + f"pinned source asset size/SHA256 mismatch for {path}: " + f"expected {expected['size']}/{expected['sha256']}, got {actual_size}/{actual_hash}" + ) + verified[relative] = actual_hash + return verified + + +def validate_qwen_vlm_destination_names( + base_assets: Path, + qwen_entry: Mapping[str, Any], + records: Sequence[TensorRecord], + *, + backbone: str, +) -> None: + """Bind Qwen2.5 staged tensor names to the pinned canonical HF weight index.""" + if backbone == "qwen3_vl": + return + if backbone != "qwen2_5_vl": + raise StarVLAError(f"unsupported StarVLA Qwen backbone: {backbone!r}") + + index_name = "model.safetensors.index.json" + if index_name not in qwen_entry.get("files", []): + raise StarVLAError("pinned Qwen2.5 asset has no canonical model weight index") + index_record = qwen_entry.get("file_hashes", {}).get(index_name) + if not isinstance(index_record, Mapping): + raise StarVLAError("pinned Qwen2.5 asset has no model weight index hash") + index_path = base_assets / index_name + if not index_path.is_file(): + raise StarVLAError(f"missing pinned Qwen2.5 model weight index: {index_path}") + if ( + index_path.stat().st_size != index_record.get("size") + or sha256_file(index_path) != index_record.get("sha256") + ): + raise StarVLAError( + f"pinned Qwen2.5 model weight index size/SHA256 mismatch: {index_path}" + ) + try: + index = json.loads(index_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise StarVLAError( + f"failed to load pinned Qwen2.5 model weight index {index_path}: {exc}" + ) from exc + weight_map = index.get("weight_map") if isinstance(index, dict) else None + if ( + not isinstance(weight_map, dict) + or not weight_map + or any( + not isinstance(name, str) + or not name + or not isinstance(shard, str) + or not shard + for name, shard in weight_map.items() + ) + ): + raise StarVLAError(f"invalid pinned Qwen2.5 model weight index: {index_path}") + + if any(record.component != "vlm" for record in records): + raise StarVLAError("Qwen VLM destination validation received a non-VLM tensor") + canonical_names = set(weight_map) + staged_backbone_names = { + record.destination_name for record in records if record.role != "lm_head" + } + staged_lm_head_names = { + record.destination_name for record in records if record.role == "lm_head" + } + if staged_lm_head_names != {"lm_head.weight"}: + raise StarVLAError( + "Qwen2.5 staged LM head tensor set mismatch: " + f"expected ['lm_head.weight'], got {sorted(staged_lm_head_names)}" + ) + if staged_backbone_names != canonical_names: + raise StarVLAError( + "Qwen2.5 staged backbone tensor names do not match the pinned canonical " + f"{len(canonical_names)}-tensor HF index; " + f"missing={sorted(canonical_names - staged_backbone_names)[:8]}, " + f"unexpected={sorted(staged_backbone_names - canonical_names)[:8]}" + ) + if len(records) != len(canonical_names) + 1: + raise StarVLAError( + "Qwen2.5 staged VLM tensor count does not equal the canonical HF index " + "plus the checkpoint LM head" + ) + + +def official_bundle_uuid(variant_entry: Mapping[str, Any], catalog: Mapping[str, Any]) -> str: + """Derive the bundle identity from every source that can change runtime semantics.""" + qwen_asset_name, qwen_entry = get_qwen_asset(catalog, variant_entry) + qwen_hashes = staged_qwen_asset_hashes(qwen_entry) + policy_hashes = { + relative: record["sha256"] for relative, record in variant_entry["file_hashes"].items() + } + provenance = { + "schema_version": 1, + "framework": variant_entry["framework"], + "policy": { + "repo_id": variant_entry["repo_id"], + "revision": variant_entry["revision"], + "checkpoint_sha256": variant_entry["checkpoint"]["sha256"], + "asset_sha256": policy_hashes, + }, + "qwen": { + "repo_id": qwen_entry["repo_id"], + "revision": qwen_entry["revision"], + "staged_asset_sha256": qwen_hashes, + }, + "source_revisions": { + "starvla": catalog["source_revisions"]["starvla"], + "llama_cpp": catalog["source_revisions"]["llama_cpp"], + }, + } + catalog_variant = variant_entry.get("_catalog_key", variant_entry["framework"]) + backbone = variant_entry.get("backbone", "qwen3_vl") + if ( + catalog_variant != variant_entry["framework"] + or backbone != "qwen3_vl" + or qwen_asset_name != DEFAULT_QWEN_ASSET + ): + provenance["catalog_variant"] = catalog_variant + provenance["backbone"] = backbone + provenance["qwen"]["asset"] = qwen_asset_name + canonical = json.dumps(provenance, sort_keys=True, separators=(",", ":")) + return str(uuid.uuid5(uuid.NAMESPACE_URL, f"robotcpp:starvla-bundle:{canonical}")) + + +def _flatten_config(value: Any, prefix: str = "") -> dict[str, Any]: + if not isinstance(value, Mapping): + return {prefix: value} + flattened: dict[str, Any] = {} + for key, child in value.items(): + path = f"{prefix}.{key}" if prefix else str(key) + flattened.update(_flatten_config(child, path)) + return flattened + + +def _set_effective_value( + effective: dict[str, Any], + path: str, + value: Any, + authority: str, + overrides: dict[str, dict[str, Any]], +) -> None: + owner: dict[str, Any] = effective + parts = path.split(".") + for part in parts[:-1]: + child = owner.get(part) + if child is None: + child = {} + owner[part] = child + if not isinstance(child, dict): + raise StarVLAError(f"effective config path is not an object: {path}") + owner = child + previous = owner.get(parts[-1]) + owner[parts[-1]] = value + if previous != value: + overrides[path] = {"source": previous, "effective": value, "authority": authority} + + +def resolve_effective_config( + source_dir: Path, + variant_name: str, + variant_entry: Mapping[str, Any] | None = None, +) -> dict[str, Any]: + """Load the runtime-canonical YAML and apply checkpoint-derived compatibility fixes.""" + framework_name = ( + str(variant_entry["framework"]) if variant_entry is not None else variant_name + ) + backbone = ( + str(variant_entry.get("backbone", "qwen3_vl")) + if variant_entry is not None + else "qwen3_vl" + ) + if framework_name not in {"oft", "groot", "pi", "pi_v3"}: + raise StarVLAError( + f"unsupported effective-config variant/framework: " + f"{variant_name!r}/{framework_name!r}" + ) + if backbone not in SUPPORTED_BACKBONES: + raise StarVLAError(f"unsupported effective-config backbone: {backbone!r}") + yaml_path = source_dir / "config.yaml" + if not yaml_path.is_file(): + raise StarVLAError(f"missing canonical StarVLA config: {yaml_path}") + try: + import yaml + except ImportError as exc: + raise StarVLAError("PyYAML is required to resolve the effective StarVLA config") from exc + try: + canonical = yaml.safe_load(yaml_path.read_text(encoding="utf-8")) + except (OSError, yaml.YAMLError) as exc: + raise StarVLAError(f"failed to load canonical StarVLA config {yaml_path}: {exc}") from exc + if not isinstance(canonical, dict): + raise StarVLAError(f"expected an object in canonical StarVLA config {yaml_path}") + + candidate_conflicts: dict[str, dict[str, Any]] = {} + json_path = source_dir / "config.json" + if json_path.is_file(): + try: + json_config = json.loads(json_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise StarVLAError(f"failed to load StarVLA config mirror {json_path}: {exc}") from exc + canonical_flat = _flatten_config(canonical) + json_flat = _flatten_config(json_config) + if set(json_flat) != set(canonical_flat): + raise StarVLAError(f"config.json and canonical config.yaml have different keys in {source_dir}") + conflicts = { + path: {"config.yaml": canonical_flat[path], "config.json": json_flat[path]} + for path in canonical_flat + if canonical_flat[path] != json_flat[path] + } + allowed = {"framework.qwenvl.base_vlm"} if framework_name == "groot" else set() + unexpected = set(conflicts) - allowed + if unexpected: + raise StarVLAError( + f"config.json and canonical config.yaml disagree at unsupported paths in {source_dir}: " + f"{sorted(unexpected)}" + ) + candidate_conflicts.update(conflicts) + elif backbone == "qwen3_vl" and framework_name in {"oft", "groot"}: + raise StarVLAError(f"missing StarVLA config mirror: {json_path}") + + if framework_name == "pi_v3": + full_path = source_dir / "config.full.yaml" + if not full_path.is_file(): + raise StarVLAError(f"missing PI_v3 full config candidate: {full_path}") + try: + full_config = yaml.safe_load(full_path.read_text(encoding="utf-8")) + except (OSError, yaml.YAMLError) as exc: + raise StarVLAError(f"failed to load PI_v3 full config {full_path}: {exc}") from exc + if not isinstance(full_config, dict): + raise StarVLAError(f"expected an object in PI_v3 full config {full_path}") + canonical_flat = _flatten_config(canonical) + full_flat = _flatten_config(full_config) + conflicts = { + path: {"config.yaml": canonical_flat[path], "config.full.yaml": full_flat[path]} + for path in set(canonical_flat) & set(full_flat) + if canonical_flat[path] != full_flat[path] + } + allowed = {"framework.action_model.diffusion_model_cfg.interleave_self_attention"} + unexpected = set(conflicts) - allowed + if unexpected: + raise StarVLAError( + f"PI_v3 config.full.yaml disagrees with canonical config.yaml at unsupported paths: " + f"{sorted(unexpected)}" + ) + candidate_conflicts.update(conflicts) + + effective = deepcopy(canonical) + overrides: dict[str, dict[str, Any]] = {} + qwen_hidden_dim = 2048 if backbone == "qwen2_5_vl" else 2560 + if framework_name == "oft": + _set_effective_value( + effective, + "framework.qwenvl.vl_hidden_dim", + qwen_hidden_dim, + "checkpoint_tensor_shape", + overrides, + ) + _set_effective_value( + effective, + "framework.action_model.action_hidden_dim", + qwen_hidden_dim, + "checkpoint_tensor_shape", + overrides, + ) + _set_effective_value( + effective, + "framework.action_model.action_model_type", + "MLP", + "pinned_starvla_qwenoft_factory_and_checkpoint_topology", + overrides, + ) + elif framework_name == "groot": + _set_effective_value( + effective, + "framework.qwenvl.vl_hidden_dim", + qwen_hidden_dim, + "checkpoint_tensor_shape", + overrides, + ) + _set_effective_value( + effective, + "framework.action_model.diffusion_model_cfg.cross_attention_dim", + qwen_hidden_dim, + "pinned_starvla_qwengroot_runtime_and_checkpoint_tensor_shape", + overrides, + ) + for path, value in ( + ("framework.action_model.diffusion_model_cfg.input_embedding_dim", 768), + ("framework.action_model.diffusion_model_cfg.attention_head_dim", 64), + ("framework.action_model.diffusion_model_cfg.num_attention_heads", 12), + ): + _set_effective_value( + effective, path, value, "pinned_starvla_dit_b_definition", overrides + ) + elif framework_name == "pi_v3": + for path, value in ( + ("framework.qwenvl.vl_hidden_dim", 2560), + ("framework.qwenvl.num_vl_layers", 36), + ("framework.action_model.action_model_type", "LayerwiseFM"), + ("framework.action_model.diffusion_model_cfg.action_dit_hidden_dim", 1024), + ("framework.action_model.diffusion_model_cfg.input_embedding_dim", 1024), + ("framework.action_model.diffusion_model_cfg.cross_attention_dim", 1024), + ("framework.action_model.diffusion_model_cfg.attention_head_dim", 64), + ("framework.action_model.diffusion_model_cfg.num_attention_heads", 16), + ("framework.action_model.diffusion_model_cfg.num_layers", 36), + ("framework.action_model.diffusion_model_cfg.interleave_self_attention", False), + ("framework.action_model.diffusion_model_cfg.use_canonical_forward", True), + ): + _set_effective_value( + effective, + path, + value, + "pinned_starvla_qwenpi_v3_runtime_and_released_checkpoint_config", + overrides, + ) + elif framework_name == "pi": + for path, value in ( + ("framework.qwenvl.vl_hidden_dim", qwen_hidden_dim), + ("framework.action_model.hidden_size", qwen_hidden_dim), + ( + "framework.action_model.diffusion_model_cfg.input_embedding_dim", + qwen_hidden_dim, + ), + ( + "framework.action_model.diffusion_model_cfg.cross_attention_dim", + qwen_hidden_dim, + ), + ("framework.action_model.diffusion_model_cfg.attention_head_dim", 64), + ( + "framework.action_model.diffusion_model_cfg.num_attention_heads", + qwen_hidden_dim // 64, + ), + ("framework.action_model.diffusion_model_cfg.use_canonical_forward", False), + ): + _set_effective_value( + effective, + path, + value, + "pinned_starvla_qwenpi_runtime_and_checkpoint_tensor_shape", + overrides, + ) + else: + raise AssertionError(f"unhandled effective-config framework: {framework_name}") + + _set_effective_value(effective, "framework.action_model.action_horizon", 16, "released_checkpoint_contract", overrides) + _set_effective_value(effective, "version_id", "0.21", "pinned_starvla_config_compat", overrides) + + inactive_fields = [] + if framework_name == "oft": + inactive_fields = [ + "framework.action_model.diffusion_model_cfg", + "framework.action_model.hidden_size", + "framework.action_model.state_dim", + ] + elif framework_name == "groot": + inactive_fields = ["framework.action_model.action_hidden_dim"] + effective["_robotcpp_effective_config"] = { + "schema_version": 1, + "variant": variant_name, + "framework": framework_name, + "backbone": backbone, + "canonical_source": "config.yaml", + "candidate_conflicts": candidate_conflicts, + "overrides": overrides, + "inactive_fields": inactive_fields, + } + return effective + + +def load_checkpoint_state(path: Path) -> dict[str, Any]: + try: + import torch + except ImportError as exc: + raise StarVLAError("PyTorch is required to inspect a StarVLA checkpoint") from exc + + try: + raw = torch.load(path, map_location="cpu", mmap=True, weights_only=True) + except Exception as exc: + raise StarVLAError(f"failed to load checkpoint {path} with weights_only=True: {exc}") from exc + + if isinstance(raw, Mapping) and raw and all(isinstance(key, str) and torch.is_tensor(value) for key, value in raw.items()): + return dict(raw) + + if isinstance(raw, Mapping): + for wrapper_key in ("state_dict", "model"): + candidate = raw.get(wrapper_key) + if isinstance(candidate, Mapping) and candidate and all( + isinstance(key, str) and torch.is_tensor(value) for key, value in candidate.items() + ): + unknown_wrappers = set(raw) - {wrapper_key} + if unknown_wrappers: + raise StarVLAError( + f"checkpoint wrapper {wrapper_key!r} has unrecognized sibling keys: {sorted(unknown_wrappers)}" + ) + return dict(candidate) + + raise StarVLAError("checkpoint must be a non-empty flat tensor state_dict or a known single-key wrapper") + + +def classify_tensor(name: str, variant_entry: Mapping[str, Any]) -> tuple[str, str, str]: + backbone = str(variant_entry.get("backbone", "qwen3_vl")) + destination_prefixes = VLM_DESTINATION_PREFIXES.get(backbone) + if destination_prefixes is None: + raise StarVLAError(f"unsupported StarVLA Qwen backbone: {backbone!r}") + + for source_prefix, role in VLM_SOURCE_RULES: + if name.startswith(source_prefix): + suffix = name[len(source_prefix) :] + if not suffix: + break + return "vlm", role, destination_prefixes[role] + suffix + + for prefix in variant_entry.get("policy_prefixes", []): + if name.startswith(prefix) and len(name) > len(prefix): + return "policy", "policy", name + if name in variant_entry.get("policy_tensors", []): + return "policy", "policy", name + + raise StarVLAError(f"unrecognized tensor for {variant_entry.get('framework')}: {name}") + + +def _storage_key(tensor: Any) -> tuple[int, int] | None: + try: + storage = tensor.untyped_storage() + return int(storage.data_ptr()), int(storage.nbytes()) + except Exception: + return None + + +def build_inventory( + state_dict: Mapping[str, Any], + variant_entry: Mapping[str, Any], + *, + enforce_expected: bool = True, +) -> list[TensorRecord]: + provisional: list[tuple[TensorRecord, tuple[int, int] | None]] = [] + destinations: dict[str, str] = {} + aliases: dict[tuple[int, int], list[str]] = defaultdict(list) + + for source_name in sorted(state_dict): + tensor = state_dict[source_name] + if not tensor.is_contiguous(): + raise StarVLAError(f"non-contiguous source tensor is not supported: {source_name}") + component, role, destination_name = classify_tensor(source_name, variant_entry) + previous = destinations.get(destination_name) + if previous is not None: + raise StarVLAError( + f"duplicate destination tensor {destination_name!r}: source keys {previous!r} and {source_name!r}" + ) + destinations[destination_name] = source_name + + storage_key = _storage_key(tensor) + if storage_key is not None: + aliases[storage_key].append(source_name) + record = TensorRecord( + source_name=source_name, + destination_name=destination_name, + component=component, + role=role, + shape=[int(dim) for dim in tensor.shape], + dtype=str(tensor.dtype).removeprefix("torch."), + numel=int(tensor.numel()), + nbytes=int(tensor.numel() * tensor.element_size()), + storage_offset=int(tensor.storage_offset()), + ) + provisional.append((record, storage_key)) + + alias_names: dict[tuple[int, int], str] = {} + alias_index = 0 + for storage_key, source_names in sorted(aliases.items(), key=lambda item: min(item[1])): + if len(source_names) > 1: + alias_names[storage_key] = f"alias_{alias_index:04d}" + alias_index += 1 + + records = [ + TensorRecord(**{**record.to_json(), "storage_alias": alias_names.get(storage_key)}) + for record, storage_key in provisional + ] + if enforce_expected: + validate_expected_inventory(records, variant_entry) + return records + + +def inventory_summary(records: list[TensorRecord]) -> dict[str, Any]: + counts = Counter(record.component for record in records) + roles = Counter(record.role for record in records) + numel = Counter() + nbytes = Counter() + dtypes = Counter() + for record in records: + numel[record.component] += record.numel + nbytes[record.component] += record.nbytes + dtypes[record.dtype] += 1 + return { + "total_tensors": len(records), + "vlm_tensors": counts["vlm"], + "policy_tensors": counts["policy"], + "visual_tensors": roles["visual"], + "text_tensors": roles["text"], + "lm_head_tensors": roles["lm_head"], + "total_numel": sum(record.numel for record in records), + "vlm_numel": numel["vlm"], + "policy_numel": numel["policy"], + "total_nbytes": sum(record.nbytes for record in records), + "vlm_nbytes": nbytes["vlm"], + "policy_nbytes": nbytes["policy"], + "dtypes": dict(sorted(dtypes.items())), + "storage_alias_groups": len({record.storage_alias for record in records if record.storage_alias}), + } + + +def validate_expected_inventory(records: list[TensorRecord], variant_entry: Mapping[str, Any]) -> None: + expected = variant_entry.get("expected") + if not expected: + raise StarVLAError(f"variant {variant_entry.get('framework')!r} has no expected checkpoint inventory") + actual = inventory_summary(records) + mismatches = [] + for key, expected_value in expected.items(): + actual_value = actual.get(key) + if actual_value != expected_value: + mismatches.append(f"{key}: expected {expected_value}, got {actual_value}") + if mismatches: + raise StarVLAError("checkpoint inventory mismatch: " + "; ".join(mismatches)) + + by_destination = {record.destination_name: record for record in records} + shape_mismatches = [] + for name, expected_shape in variant_entry.get("required_shapes", {}).items(): + record = by_destination.get(name) + if record is None: + shape_mismatches.append(f"{name}: missing") + elif record.shape != expected_shape: + shape_mismatches.append(f"{name}: expected {expected_shape}, got {record.shape}") + if shape_mismatches: + raise StarVLAError("checkpoint required-shape mismatch: " + "; ".join(shape_mismatches)) + + +def validate_official_surgery_manifest( + manifest: Mapping[str, Any], + variant_entry: Mapping[str, Any], + catalog: Mapping[str, Any], +) -> None: + """Require a surgery manifest to describe the pinned official checkpoint exactly.""" + checkpoint = variant_entry.get("checkpoint") + if checkpoint is None: + raise StarVLAError(f"variant {variant_entry.get('framework')!r} has no official checkpoint") + + expected_top_level = { + "schema_version": 1, + "variant": variant_entry.get("_catalog_key", variant_entry["framework"]), + "model_type": variant_entry["model_type"], + } + _, qwen_entry = get_qwen_asset(catalog, variant_entry) + source = manifest.get("source") + if not isinstance(source, Mapping): + raise StarVLAError("surgery manifest has no source object") + expected_source = { + "repo_id": variant_entry["repo_id"], + "revision": variant_entry["revision"], + "checkpoint_size": checkpoint["size"], + "checkpoint_sha256": checkpoint["sha256"], + "starvla_revision": catalog["source_revisions"]["starvla"], + "llama_cpp_revision": catalog["source_revisions"]["llama_cpp"], + "qwen_repo_id": qwen_entry["repo_id"], + "qwen_revision": qwen_entry["revision"], + } + mismatches = [] + for key, expected in expected_top_level.items(): + if manifest.get(key) != expected: + mismatches.append(f"{key}: expected {expected!r}, got {manifest.get(key)!r}") + for key, expected in expected_source.items(): + if source.get(key) != expected: + mismatches.append(f"source.{key}: expected {expected!r}, got {source.get(key)!r}") + + inventory = manifest.get("inventory") + if not isinstance(inventory, Mapping): + mismatches.append("inventory: missing or not an object") + else: + for key, expected in variant_entry.get("expected", {}).items(): + if inventory.get(key) != expected: + mismatches.append(f"inventory.{key}: expected {expected!r}, got {inventory.get(key)!r}") + + expected_uuid = official_bundle_uuid(variant_entry, catalog) + if manifest.get("bundle_uuid") != expected_uuid: + mismatches.append(f"bundle_uuid: expected {expected_uuid!r}, got {manifest.get('bundle_uuid')!r}") + + qwen_expected = staged_qwen_asset_hashes(qwen_entry) + policy_expected = { + relative: record["sha256"] for relative, record in variant_entry["file_hashes"].items() + } + if manifest.get("qwen_assets") != qwen_expected: + mismatches.append("qwen_assets: staged hashes do not match the pinned Qwen assets") + if manifest.get("policy_assets") != policy_expected: + mismatches.append("policy_assets: staged hashes do not match the pinned policy assets") + + tensors = manifest.get("tensors") + expected_tensor_count = int(variant_entry.get("expected", {}).get("total_tensors", -1)) + if not isinstance(tensors, list) or len(tensors) != expected_tensor_count: + mismatches.append( + f"tensors: expected a {expected_tensor_count}-record inventory, " + f"got {len(tensors) if isinstance(tensors, list) else 'missing'}" + ) + else: + required_record_keys = set(TensorRecord.__dataclass_fields__) + source_names = set() + destination_names = set() + for index, record in enumerate(tensors): + if not isinstance(record, Mapping) or set(record) != required_record_keys: + mismatches.append(f"tensors[{index}]: invalid tensor record schema") + break + source_names.add(record["source_name"]) + destination_names.add(record["destination_name"]) + if len(source_names) != expected_tensor_count or len(destination_names) != expected_tensor_count: + mismatches.append("tensors: source and destination names must be unique") + + for field in ("vlm_output", "policy_output"): + output = manifest.get(field) + if not isinstance(output, Mapping): + mismatches.append(f"{field}: missing or not an object") + continue + if not isinstance(output.get("index"), str) or not isinstance(output.get("shards"), list): + mismatches.append(f"{field}: invalid index/shard records") + if not isinstance(output.get("index_size"), int) or not isinstance(output.get("index_sha256"), str): + mismatches.append(f"{field}: missing index size/SHA256") + + effective_config = manifest.get("effective_config") + if not isinstance(effective_config, Mapping): + mismatches.append("effective_config: missing or not an object") + elif ( + effective_config.get("path") != "effective_config.json" + or not isinstance(effective_config.get("size"), int) + or not isinstance(effective_config.get("sha256"), str) + ): + mismatches.append("effective_config: invalid path/size/SHA256 record") + if mismatches: + raise StarVLAError("non-official or inconsistent surgery manifest: " + "; ".join(mismatches)) + + +def verify_staged_assets(root: Path, assets: Mapping[str, Any], *, component: str) -> None: + if not isinstance(assets, Mapping) or not assets: + raise StarVLAError(f"surgery manifest has no {component} asset hashes") + for relative, expected_hash in sorted(assets.items()): + relative_path = _safe_relative_path(relative, field=f"{component} assets") + path = root / relative_path + if not path.is_file(): + raise StarVLAError(f"missing staged {component} asset: {path}") + actual_hash = sha256_file(path) + if actual_hash != expected_hash: + raise StarVLAError( + f"staged {component} asset SHA256 mismatch for {path}: expected {expected_hash}, got {actual_hash}" + ) + + +def verify_staged_shards(root: Path, output: Mapping[str, Any], *, component: str) -> dict[str, Any]: + if not isinstance(output, Mapping): + raise StarVLAError(f"surgery manifest has no {component} output object") + index_relative = _safe_relative_path(output.get("index"), field=f"{component} index") + index_path = root / index_relative + if not index_path.is_file(): + raise StarVLAError(f"missing staged {component} index: {index_path}") + expected_index_size = output.get("index_size") + expected_index_hash = output.get("index_sha256") + if index_path.stat().st_size != expected_index_size or sha256_file(index_path) != expected_index_hash: + raise StarVLAError(f"staged {component} index size/SHA256 mismatch: {index_path}") + try: + index = json.loads(index_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise StarVLAError(f"failed to load staged {component} index {index_path}: {exc}") from exc + weight_map = index.get("weight_map") if isinstance(index, dict) else None + if not isinstance(weight_map, dict) or not weight_map: + raise StarVLAError(f"invalid or empty staged {component} weight_map: {index_path}") + + shard_records = output.get("shards") + if not isinstance(shard_records, list) or not shard_records: + raise StarVLAError(f"surgery manifest has no {component} shard records") + manifest_shards = set() + tensor_count = 0 + for record in shard_records: + if not isinstance(record, Mapping): + raise StarVLAError(f"invalid {component} shard record: {record!r}") + relative_path = _safe_relative_path(record.get("path"), field=f"{component} shard") + relative = relative_path.as_posix() + if relative in manifest_shards: + raise StarVLAError(f"duplicate {component} shard record: {relative}") + manifest_shards.add(relative) + path = root / relative_path + if not path.is_file(): + raise StarVLAError(f"missing staged {component} shard: {path}") + actual_size = path.stat().st_size + actual_hash = sha256_file(path) + if actual_size != record.get("size") or actual_hash != record.get("sha256"): + raise StarVLAError(f"staged {component} shard size/SHA256 mismatch: {path}") + tensor_count += int(record.get("tensor_count", -1)) + + indexed_shards = {str(value) for value in weight_map.values()} + if indexed_shards != manifest_shards: + raise StarVLAError( + f"staged {component} index/manifest shard mismatch: index={sorted(indexed_shards)}, " + f"manifest={sorted(manifest_shards)}" + ) + if tensor_count != len(weight_map): + raise StarVLAError( + f"staged {component} tensor count mismatch: manifest={tensor_count}, index={len(weight_map)}" + ) + return index + + +def _verify_staged_component( + root: Path, + output: Mapping[str, Any], + state_dict: Mapping[str, Any], + source_records: list[TensorRecord], + *, + component: str, +) -> None: + index = verify_staged_shards(root, output, component=component) + weight_map = index["weight_map"] + records = [record for record in source_records if record.component == component] + expected_names = {record.destination_name for record in records} + if set(weight_map) != expected_names: + raise StarVLAError( + f"staged {component} tensor set does not match the official checkpoint: " + f"expected {len(expected_names)}, got {len(weight_map)}" + ) + + try: + import torch + from safetensors import safe_open + except ImportError as exc: + raise StarVLAError("PyTorch and safetensors are required for staged tensor verification") from exc + + by_shard: dict[str, list[TensorRecord]] = defaultdict(list) + for record in records: + by_shard[str(weight_map[record.destination_name])].append(record) + for shard, shard_records in sorted(by_shard.items()): + shard_path = root / _safe_relative_path(shard, field=f"{component} shard index") + with safe_open(shard_path, framework="pt", device="cpu") as handle: + names = {record.destination_name for record in shard_records} + if set(handle.keys()) != names: + raise StarVLAError(f"staged {component} shard/index key mismatch: {shard_path}") + for record in shard_records: + staged = handle.get_tensor(record.destination_name) + original = state_dict[record.source_name].detach().cpu() + if staged.dtype != original.dtype or list(staged.shape) != list(original.shape): + raise StarVLAError( + f"staged tensor dtype/shape mismatch for {record.destination_name}: " + f"expected {original.dtype}/{list(original.shape)}, " + f"got {staged.dtype}/{list(staged.shape)}" + ) + if not torch.equal(staged, original): + raise StarVLAError( + f"staged tensor content does not match the official checkpoint: {record.destination_name}" + ) + del staged + + +def verify_staged_components_against_checkpoint( + components: Mapping[str, tuple[Path, Mapping[str, Any]]], + manifest: Mapping[str, Any], + variant_entry: Mapping[str, Any], +) -> None: + """Bind one or more staged components to a single load of the pinned checkpoint.""" + if not components or not set(components).issubset({"vlm", "policy"}): + raise StarVLAError(f"invalid staged tensor components: {sorted(components)}") + source = manifest.get("source") + if not isinstance(source, Mapping) or not isinstance(source.get("checkpoint"), str): + raise StarVLAError("surgery manifest has no source checkpoint path") + checkpoint_path = Path(source["checkpoint"]) + verify_checkpoint_file(checkpoint_path, variant_entry) + + state_dict = load_checkpoint_state(checkpoint_path) + source_records = build_inventory(state_dict, variant_entry, enforce_expected=True) + manifest_records = manifest.get("tensors") + expected_manifest = [record.to_json() for record in source_records] + if manifest_records != expected_manifest: + raise StarVLAError("surgery tensor inventory does not match the verified official checkpoint") + for component, (root, output) in components.items(): + _verify_staged_component( + root, + output, + state_dict, + source_records, + component=component, + ) + del state_dict + + +def verify_staged_tensors_against_checkpoint( + root: Path, + output: Mapping[str, Any], + manifest: Mapping[str, Any], + variant_entry: Mapping[str, Any], + *, + component: str, +) -> None: + """Bind one staged component to the pinned checkpoint.""" + verify_staged_components_against_checkpoint( + {component: (root, output)}, + manifest, + variant_entry, + ) + + +def create_output_temporary(path: Path) -> tuple[int, Path]: + """Create a same-directory temporary file with normal umask-derived permissions.""" + path.parent.mkdir(parents=True, exist_ok=True) + for _ in range(100): + temporary = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp") + try: + descriptor = os.open( + temporary, + os.O_CREAT | os.O_EXCL | os.O_WRONLY, + 0o666, + ) + except FileExistsError: + continue + return descriptor, temporary + raise StarVLAError(f"failed to allocate a temporary output beside {path}") + + +def atomic_write_json(path: Path, value: Any, *, overwrite: bool = True) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + payload = json.dumps(value, indent=2, sort_keys=True) + "\n" + descriptor, temporary = create_output_temporary(path) + try: + with os.fdopen(descriptor, "w", encoding="utf-8") as handle: + handle.write(payload) + handle.flush() + os.fsync(handle.fileno()) + if overwrite: + os.replace(temporary, path) + else: + try: + os.link(temporary, path) + except FileExistsError as exc: + raise StarVLAError(f"refusing to overwrite existing output: {path}") from exc + temporary.unlink() + finally: + temporary.unlink(missing_ok=True) diff --git a/tools/hf2gguf/starvla/starvla_surgery.py b/tools/hf2gguf/starvla/starvla_surgery.py new file mode 100755 index 0000000..54831ed --- /dev/null +++ b/tools/hf2gguf/starvla/starvla_surgery.py @@ -0,0 +1,368 @@ +#!/usr/bin/env python3 +"""Split a StarVLA checkpoint into HF-compatible Qwen3-VL and policy staging.""" + +from __future__ import annotations + +import argparse +import json +import os +import re +import shutil +import sys +from pathlib import Path +from typing import Any, Mapping + +from starvla_checkpoint import ( + DEFAULT_CATALOG, + StarVLAError, + TensorRecord, + atomic_write_json, + build_inventory, + get_qwen_asset, + get_variant, + inventory_summary, + load_catalog, + load_checkpoint_state, + official_bundle_uuid, + resolve_effective_config, + sha256_file, + staged_qwen_asset_hashes, + validate_qwen_vlm_destination_names, + verify_catalog_files, + verify_checkpoint_file, +) + + +def parse_size(value: str) -> int: + match = re.fullmatch(r"([1-9][0-9]*)([KMG]?)", value.strip().upper()) + if not match: + raise argparse.ArgumentTypeError("size must be an integer optionally followed by K, M, or G") + amount = int(match.group(1)) + multiplier = {"": 1, "K": 1024, "M": 1024**2, "G": 1024**3}[match.group(2)] + return amount * multiplier + + +def plan_shards(records: list[TensorRecord], max_shard_size: int) -> list[list[TensorRecord]]: + if max_shard_size <= 0: + raise StarVLAError("max shard size must be positive") + shards: list[list[TensorRecord]] = [] + current: list[TensorRecord] = [] + current_size = 0 + for record in sorted(records, key=lambda item: item.destination_name): + if current and current_size + record.nbytes > max_shard_size: + shards.append(current) + current = [] + current_size = 0 + current.append(record) + current_size += record.nbytes + if current: + shards.append(current) + return shards + + +def _prepare_tensor(tensor: Any, *, clone: bool) -> Any: + prepared = tensor.detach().cpu() + if clone: + prepared = prepared.clone() + elif not prepared.is_contiguous(): + prepared = prepared.contiguous() + return prepared + + +def write_safetensor_shards( + output_dir: Path, + prefix: str, + index_name: str, + records: list[TensorRecord], + state_dict: Mapping[str, Any], + max_shard_size: int, +) -> dict[str, Any]: + try: + from safetensors import safe_open + from safetensors.torch import save_file + except ImportError as exc: + raise StarVLAError("safetensors is required for StarVLA surgery") from exc + + output_dir.mkdir(parents=True, exist_ok=True) + shards = plan_shards(records, max_shard_size) + weight_map: dict[str, str] = {} + shard_manifest = [] + alias_seen: set[str] = set() + + for index, shard_records in enumerate(shards, start=1): + filename = f"{prefix}-{index:05d}-of-{len(shards):05d}.safetensors" + path = output_dir / filename + temporary = output_dir / f".{filename}.tmp" + tensors = {} + for record in shard_records: + clone = bool(record.storage_alias and record.storage_alias in alias_seen) + tensors[record.destination_name] = _prepare_tensor(state_dict[record.source_name], clone=clone) + if record.storage_alias: + alias_seen.add(record.storage_alias) + weight_map[record.destination_name] = filename + + save_file(tensors, temporary, metadata={"format": "pt"}) + os.replace(temporary, path) + with safe_open(path, framework="pt", device="cpu") as handle: + actual_names = set(handle.keys()) + expected_names = set(tensors) + if actual_names != expected_names: + raise StarVLAError(f"safetensors key mismatch after writing {path}") + for record in shard_records: + actual_shape = [int(dim) for dim in handle.get_slice(record.destination_name).get_shape()] + if actual_shape != record.shape: + raise StarVLAError( + f"safetensors shape mismatch for {record.destination_name}: " + f"expected {record.shape}, got {actual_shape}" + ) + shard_manifest.append( + { + "path": filename, + "size": path.stat().st_size, + "sha256": sha256_file(path), + "tensor_count": len(shard_records), + } + ) + del tensors + + expected_destinations = {record.destination_name for record in records} + if set(weight_map) != expected_destinations: + raise StarVLAError("safetensors weight map does not cover every destination tensor exactly once") + index = { + "metadata": {"total_size": sum(record.nbytes for record in records)}, + "weight_map": dict(sorted(weight_map.items())), + } + index_path = output_dir / index_name + atomic_write_json(index_path, index) + return { + "index": index_name, + "index_size": index_path.stat().st_size, + "index_sha256": sha256_file(index_path), + "shards": shard_manifest, + } + + +def copy_qwen_assets(base_assets: Path, hf_dir: Path, asset_entry: Mapping[str, Any]) -> dict[str, str]: + verify_catalog_files(base_assets, asset_entry) + expected_staged_assets = staged_qwen_asset_hashes(asset_entry) + copied = {} + for relative in expected_staged_assets: + source = base_assets / relative + if not source.is_file(): + raise StarVLAError(f"missing pinned Qwen asset: {source}") + destination = hf_dir / relative + destination.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(source, destination) + copied[relative] = sha256_file(destination) + + config_path = hf_dir / "config.json" + config = json.loads(config_path.read_text(encoding="utf-8")) + config["tie_word_embeddings"] = False + text_config = config.get("text_config") + if isinstance(text_config, dict): + text_config["tie_word_embeddings"] = False + elif config.get("model_type") != "qwen2_5_vl": + raise StarVLAError("Qwen-VL config.json has no text_config object") + atomic_write_json(config_path, config) + copied["config.json"] = sha256_file(config_path) + if copied != expected_staged_assets: + raise StarVLAError( + "staged Qwen asset hashes do not match the catalog overrides" + ) + return copied + + +def copy_policy_assets(source_dir: Path, policy_dir: Path, variant_entry: Mapping[str, Any]) -> dict[str, str]: + verify_catalog_files(source_dir, variant_entry) + copied = {} + for relative in variant_entry.get("files", []): + source = source_dir / relative + if not source.is_file(): + raise StarVLAError(f"missing pinned StarVLA policy asset: {source}") + destination = policy_dir / relative + destination.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(source, destination) + copied[relative] = sha256_file(destination) + return copied + + +def _run_surgery_in_owned_directory( + checkpoint: Path, + source_dir: Path, + base_assets: Path, + output_dir: Path, + variant_name: str, + catalog_path: Path, + max_shard_size: int, + *, + verify_hash: bool, + enforce_expected: bool, +) -> dict[str, Any]: + catalog = load_catalog(catalog_path) + variant = get_variant(catalog, variant_name) + if variant.get("checkpoint") is None: + raise StarVLAError(f"variant {variant_name!r} has no official policy checkpoint to split") + if verify_hash: + verify_checkpoint_file(checkpoint, variant) + if output_dir.exists() and any(output_dir.iterdir()): + raise StarVLAError(f"output directory is not empty: {output_dir}") + + state_dict = load_checkpoint_state(checkpoint) + records = build_inventory(state_dict, variant, enforce_expected=enforce_expected) + vlm_records = [record for record in records if record.component == "vlm"] + policy_records = [record for record in records if record.component == "policy"] + if len(vlm_records) + len(policy_records) != len(records): + raise StarVLAError("source tensor set is not the disjoint union of VLM and policy tensors") + + hf_dir = output_dir / "hf" + policy_dir = output_dir / "policy" + hf_dir.mkdir(parents=True, exist_ok=True) + policy_dir.mkdir(parents=True, exist_ok=True) + + qwen_asset_name, qwen_asset_entry = get_qwen_asset(catalog, variant) + validate_qwen_vlm_destination_names( + base_assets, + qwen_asset_entry, + vlm_records, + backbone=str(variant.get("backbone", "qwen3_vl")), + ) + qwen_assets = copy_qwen_assets(base_assets, hf_dir, qwen_asset_entry) + policy_assets = copy_policy_assets(source_dir, policy_dir, variant) + effective_config_path = policy_dir / "effective_config.json" + atomic_write_json( + effective_config_path, + resolve_effective_config(source_dir, variant_name, variant), + ) + effective_config = { + "path": effective_config_path.name, + "size": effective_config_path.stat().st_size, + "sha256": sha256_file(effective_config_path), + } + vlm_output = write_safetensor_shards( + hf_dir, + "model", + "model.safetensors.index.json", + vlm_records, + state_dict, + max_shard_size, + ) + policy_output = write_safetensor_shards( + policy_dir, + "policy", + "policy.safetensors.index.json", + policy_records, + state_dict, + max_shard_size, + ) + + checkpoint_sha256 = sha256_file(checkpoint) if not verify_hash else str(variant["checkpoint"]["sha256"]) + bundle_uuid = official_bundle_uuid(variant, catalog) + manifest = { + "schema_version": 1, + "variant": variant_name, + "framework": variant["framework"], + "backbone": variant.get("backbone", "qwen3_vl"), + "model_type": variant["model_type"], + "bundle_uuid": bundle_uuid, + "source": { + "repo_id": variant["repo_id"], + "revision": variant["revision"], + "checkpoint": str(checkpoint), + "checkpoint_size": checkpoint.stat().st_size, + "checkpoint_sha256": checkpoint_sha256, + "starvla_revision": catalog["source_revisions"]["starvla"], + "llama_cpp_revision": catalog["source_revisions"]["llama_cpp"], + "qwen_repo_id": qwen_asset_entry["repo_id"], + "qwen_revision": qwen_asset_entry["revision"], + "qwen_asset": qwen_asset_name, + }, + "inventory": inventory_summary(records), + "qwen_assets": qwen_assets, + "policy_assets": policy_assets, + "effective_config": effective_config, + "vlm_output": vlm_output, + "policy_output": policy_output, + "tensors": [record.to_json() for record in records], + } + atomic_write_json(output_dir / "surgery_manifest.json", manifest, overwrite=False) + return manifest + + +def run_surgery( + checkpoint: Path, + source_dir: Path, + base_assets: Path, + output_dir: Path, + variant_name: str, + catalog_path: Path, + max_shard_size: int, + *, + verify_hash: bool, + enforce_expected: bool, +) -> dict[str, Any]: + """Own the staging directory so a failed split cannot poison a retry.""" + output_dir.parent.mkdir(parents=True, exist_ok=True) + try: + output_dir.mkdir() + except FileExistsError as exc: + raise StarVLAError(f"refusing to overwrite existing output directory: {output_dir}") from exc + + try: + return _run_surgery_in_owned_directory( + checkpoint=checkpoint, + source_dir=source_dir, + base_assets=base_assets, + output_dir=output_dir, + variant_name=variant_name, + catalog_path=catalog_path, + max_shard_size=max_shard_size, + verify_hash=verify_hash, + enforce_expected=enforce_expected, + ) + except BaseException: + shutil.rmtree(output_dir, ignore_errors=True) + raise + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("checkpoint", type=Path) + parser.add_argument( + "--variant", + required=True, + choices=("oft", "groot", "pi_v3", "qwen25_oft", "qwen25_groot", "qwen25_pi"), + ) + parser.add_argument("--source-dir", type=Path, required=True) + parser.add_argument("--base-assets", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--catalog", type=Path, default=DEFAULT_CATALOG) + parser.add_argument("--max-shard-size", type=parse_size, default=parse_size("2G")) + parser.add_argument("--skip-hash-check", action="store_true") + parser.add_argument("--allow-nonofficial-inventory", action="store_true") + return parser.parse_args() + + +def main() -> int: + args = parse_args() + try: + manifest = run_surgery( + checkpoint=args.checkpoint, + source_dir=args.source_dir, + base_assets=args.base_assets, + output_dir=args.output_dir, + variant_name=args.variant, + catalog_path=args.catalog, + max_shard_size=args.max_shard_size, + verify_hash=not args.skip_hash_check, + enforce_expected=not args.allow_nonofficial_inventory, + ) + print(json.dumps(manifest["inventory"], indent=2, sort_keys=True)) + print(f"surgery manifest: {args.output_dir / 'surgery_manifest.json'}") + return 0 + except (StarVLAError, OSError, json.JSONDecodeError) as exc: + print(f"error: {exc}", file=sys.stderr) + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/hf2gguf/starvla/validate_starvla_bundle.py b/tools/hf2gguf/starvla/validate_starvla_bundle.py new file mode 100755 index 0000000..560b27a --- /dev/null +++ b/tools/hf2gguf/starvla/validate_starvla_bundle.py @@ -0,0 +1,1224 @@ +#!/usr/bin/env python3 +"""Validate a converted StarVLA GGUF bundle and write its content manifest.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +import sys +from collections import Counter +from pathlib import Path +from typing import Any + +import numpy as np + +from convert_starvla_policy_to_gguf import ( + GROOT_BLOCK_COUNT, + GROOT_OFFICIAL_DIMENSIONS, + GROOT_OFFICIAL_DIMENSIONS_BY_BACKBONE, + GROOT_POLICY_TENSOR_COUNT, + GROOT_TENSOR_MAP, + OFT_ACTION_TOKEN_ID, + OFT_TENSOR_MAP, + PI_BLOCK_COUNT, + PI_OFFICIAL_DIMENSIONS, + PI_POLICY_TENSOR_COUNT, + PI_TENSOR_MAP, + PI_V3_BLOCK_COUNT, + PI_V3_OFFICIAL_DIMENSIONS, + PI_V3_POLICY_TENSOR_COUNT, + PI_V3_PROJECTOR_COUNT, + PI_V3_TENSOR_MAP, + QWEN3VL_DYNAMIC_IMAGE_METADATA, + build_groot_metadata, + build_oft_metadata, + build_pi_metadata, + build_pi_v3_metadata, + load_policy_tensors, + resolve_action_token_id, +) +from starvla_checkpoint import ( + DEFAULT_CATALOG, + DEFAULT_MMPROJ_DTYPE, + DEFAULT_POLICY_DTYPE, + DEFAULT_TEXT_DTYPE, + StarVLAError, + atomic_write_json, + get_variant, + load_catalog, + sha256_file, + validate_official_surgery_manifest, + verify_staged_assets, + verify_staged_components_against_checkpoint, +) + + +REPOSITORY_ROOT = Path(__file__).resolve().parents[3] +LLAMA_GGUF_PY = REPOSITORY_ROOT / "third_party" / "llama.cpp" / "gguf-py" +sys.path.insert(0, str(LLAMA_GGUF_PY)) + +try: + import gguf +except ImportError as exc: + raise SystemExit(f"error: failed to import pinned llama.cpp gguf-py: {exc}") from exc + + +def _load_json(path: Path) -> dict[str, Any]: + try: + value = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise StarVLAError(f"failed to load JSON {path}: {exc}") from exc + if not isinstance(value, dict): + raise StarVLAError(f"expected a JSON object in {path}") + return value + + +def field_value(reader: Any, key: str) -> Any: + field = reader.get_field(key) + if field is None: + raise StarVLAError(f"GGUF is missing required metadata: {key}") + return field.contents() + + +def expect_field(reader: Any, key: str, expected: Any) -> None: + actual = field_value(reader, key) + if actual != expected: + raise StarVLAError(f"GGUF metadata mismatch for {key}: expected {expected!r}, got {actual!r}") + + +def expect_sequence_field(reader: Any, key: str, expected: list[Any]) -> None: + actual = field_value(reader, key) + if not isinstance(actual, list): + raise StarVLAError(f"GGUF metadata mismatch for {key}: expected an array, got {type(actual).__name__}") + if len(actual) != len(expected): + raise StarVLAError( + f"GGUF metadata length mismatch for {key}: expected {len(expected)}, got {len(actual)}" + ) + for index, (actual_item, expected_item) in enumerate(zip(actual, expected)): + if actual_item != expected_item: + raise StarVLAError( + f"GGUF metadata mismatch for {key}[{index}]: " + f"expected {expected_item!r}, got {actual_item!r}" + ) + + +def tensor_map(reader: Any) -> dict[str, Any]: + tensors = {tensor.name: tensor for tensor in reader.tensors} + if len(tensors) != len(reader.tensors): + raise StarVLAError("GGUF contains duplicate tensor names") + return tensors + + +def expect_ggml_tensor_shape(tensors: dict[str, Any], name: str, shape: list[int]) -> None: + """Check GGUF/ggml dimensions (`ne[]` order), not NumPy/PyTorch dimensions.""" + tensor = tensors.get(name) + if tensor is None: + raise StarVLAError(f"GGUF is missing required tensor: {name}") + actual = [int(dim) for dim in tensor.shape] + if actual != shape: + raise StarVLAError(f"GGUF tensor shape mismatch for {name}: expected {shape}, got {actual}") + + +def expect_complete_tensor_map(tensors: dict[str, Any], expected: dict[str, list[int]], component: str) -> None: + actual_names = set(tensors) + expected_names = set(expected) + if actual_names != expected_names: + raise StarVLAError( + f"{component} GGUF tensor set mismatch; " + f"missing={sorted(expected_names - actual_names)}, unexpected={sorted(actual_names - expected_names)}" + ) + for name, shape in expected.items(): + expect_ggml_tensor_shape(tensors, name, shape) + + +def expected_text_tensor_map( + backbone: str = "qwen3_vl", vocab_size: int = 151936 +) -> dict[str, list[int]]: + if backbone == "qwen2_5_vl": + expected = { + "token_embd.weight": [2048, vocab_size], + "output_norm.weight": [2048], + "output.weight": [2048, vocab_size], + } + per_block = { + "attn_norm.weight": [2048], + "ffn_norm.weight": [2048], + "attn_q.weight": [2048, 2048], + "attn_q.bias": [2048], + "attn_k.weight": [2048, 256], + "attn_k.bias": [256], + "attn_v.weight": [2048, 256], + "attn_v.bias": [256], + "attn_output.weight": [2048, 2048], + "ffn_gate.weight": [2048, 11008], + "ffn_up.weight": [2048, 11008], + "ffn_down.weight": [11008, 2048], + } + for block in range(36): + for suffix, shape in per_block.items(): + expected[f"blk.{block}.{suffix}"] = shape + return expected + if backbone != "qwen3_vl": + raise StarVLAError(f"unsupported Qwen text tensor backbone: {backbone!r}") + expected = { + "token_embd.weight": [2560, 151936], + "output_norm.weight": [2560], + "output.weight": [2560, 151936], + } + per_block = { + "attn_norm.weight": [2560], + "ffn_norm.weight": [2560], + "attn_q.weight": [2560, 4096], + "attn_k.weight": [2560, 1024], + "attn_v.weight": [2560, 1024], + "attn_output.weight": [4096, 2560], + "attn_q_norm.weight": [128], + "attn_k_norm.weight": [128], + "ffn_gate.weight": [2560, 9728], + "ffn_up.weight": [2560, 9728], + "ffn_down.weight": [9728, 2560], + } + for block in range(36): + for suffix, shape in per_block.items(): + expected[f"blk.{block}.{suffix}"] = shape + return expected + + +def expected_mmproj_tensor_map( + backbone: str = "qwen3_vl", +) -> dict[str, list[int]]: + if backbone == "qwen2_5_vl": + expected = { + "v.patch_embd.weight": [14, 14, 3, 1280], + "v.patch_embd.weight.1": [14, 14, 3, 1280], + "v.post_ln.weight": [1280], + "mm.0.weight": [5120, 5120], + "mm.0.bias": [5120], + "mm.2.weight": [5120, 2048], + "mm.2.bias": [2048], + } + per_block = { + "ln1.weight": [1280], + "ln2.weight": [1280], + "attn_q.weight": [1280, 1280], + "attn_q.bias": [1280], + "attn_k.weight": [1280, 1280], + "attn_k.bias": [1280], + "attn_v.weight": [1280, 1280], + "attn_v.bias": [1280], + "attn_out.weight": [1280, 1280], + "attn_out.bias": [1280], + "ffn_gate.weight": [1280, 3420], + "ffn_gate.bias": [3420], + "ffn_up.weight": [1280, 3420], + "ffn_up.bias": [3420], + "ffn_down.weight": [3420, 1280], + "ffn_down.bias": [1280], + } + for block in range(32): + for suffix, shape in per_block.items(): + expected[f"v.blk.{block}.{suffix}"] = shape + return expected + if backbone != "qwen3_vl": + raise StarVLAError(f"unsupported Qwen mmproj tensor backbone: {backbone!r}") + expected = { + "v.position_embd.weight": [1024, 2304], + "v.patch_embd.weight": [16, 16, 3, 1024], + "v.patch_embd.weight.1": [16, 16, 3, 1024], + "v.patch_embd.bias": [1024], + "v.post_ln.weight": [1024], + "v.post_ln.bias": [1024], + "mm.0.weight": [4096, 4096], + "mm.0.bias": [4096], + "mm.2.weight": [4096, 2560], + "mm.2.bias": [2560], + } + per_block = { + "attn_out.weight": [1024, 1024], + "attn_out.bias": [1024], + "attn_qkv.weight": [1024, 3072], + "attn_qkv.bias": [3072], + "ffn_up.weight": [1024, 4096], + "ffn_up.bias": [4096], + "ffn_down.weight": [4096, 1024], + "ffn_down.bias": [1024], + "ln1.weight": [1024], + "ln1.bias": [1024], + "ln2.weight": [1024], + "ln2.bias": [1024], + } + for block in range(24): + for suffix, shape in per_block.items(): + expected[f"v.blk.{block}.{suffix}"] = shape + for layer in (5, 11, 17): + expected.update( + { + f"v.deepstack.{layer}.norm.weight": [4096], + f"v.deepstack.{layer}.norm.bias": [4096], + f"v.deepstack.{layer}.fc1.weight": [4096, 4096], + f"v.deepstack.{layer}.fc1.bias": [4096], + f"v.deepstack.{layer}.fc2.weight": [4096, 2560], + f"v.deepstack.{layer}.fc2.bias": [2560], + } + ) + return expected + + +def expected_groot_policy_tensor_map( + qwen_hidden_dim: int = 2560, +) -> dict[str, list[int]]: + """Return all GR00T GGUF shapes in ggml `ne[]` dimension order.""" + expected = { + "starvla.policy.groot.timestep.input.weight": [256, 768], + "starvla.policy.groot.timestep.input.bias": [768], + "starvla.policy.groot.timestep.output.weight": [768, 768], + "starvla.policy.groot.timestep.output.bias": [768], + } + for block in range(GROOT_BLOCK_COUNT): + attention_input_dim = qwen_hidden_dim if block % 2 == 0 else 768 + prefix = f"starvla.policy.groot.block.{block}" + expected.update( + { + f"{prefix}.ada_norm.weight": [768, 1536], + f"{prefix}.ada_norm.bias": [1536], + f"{prefix}.attention.query.weight": [768, 768], + f"{prefix}.attention.query.bias": [768], + f"{prefix}.attention.key.weight": [attention_input_dim, 768], + f"{prefix}.attention.key.bias": [768], + f"{prefix}.attention.value.weight": [attention_input_dim, 768], + f"{prefix}.attention.value.bias": [768], + f"{prefix}.attention.output.weight": [768, 768], + f"{prefix}.attention.output.bias": [768], + f"{prefix}.feed_forward.input.weight": [768, 3072], + f"{prefix}.feed_forward.input.bias": [3072], + f"{prefix}.feed_forward.output.weight": [3072, 768], + f"{prefix}.feed_forward.output.bias": [768], + } + ) + expected.update( + { + "starvla.policy.groot.output.modulation.weight": [768, 1536], + "starvla.policy.groot.output.modulation.bias": [1536], + "starvla.policy.groot.output.projection.weight": [768, 1024], + "starvla.policy.groot.output.projection.bias": [1024], + "starvla.policy.groot.action.input.weight": [7, 768], + "starvla.policy.groot.action.input.bias": [768], + "starvla.policy.groot.action.time_mix.weight": [1536, 768], + "starvla.policy.groot.action.time_mix.bias": [768], + "starvla.policy.groot.action.output.weight": [768, 768], + "starvla.policy.groot.action.output.bias": [768], + "starvla.policy.groot.velocity.input.weight": [1024, 1024], + "starvla.policy.groot.velocity.input.bias": [1024], + "starvla.policy.groot.velocity.output.weight": [1024, 7], + "starvla.policy.groot.velocity.output.bias": [7], + "starvla.policy.groot.future_tokens.weight": [768, 32], + "starvla.policy.groot.action_position.weight": [768, 1024], + } + ) + if len(expected) != GROOT_POLICY_TENSOR_COUNT: + raise AssertionError(f"internal GR00T tensor contract has {len(expected)} tensors") + return expected + + +def expected_pi_policy_tensor_map() -> dict[str, list[int]]: + """Return all legacy PI GGUF shapes in ggml `ne[]` dimension order.""" + expected = { + "starvla.policy.pi.timestep.input.weight": [256, 2048], + "starvla.policy.pi.timestep.input.bias": [2048], + "starvla.policy.pi.timestep.output.weight": [2048, 2048], + "starvla.policy.pi.timestep.output.bias": [2048], + } + for block in range(PI_BLOCK_COUNT): + prefix = f"starvla.policy.pi.block.{block}" + expected.update( + { + f"{prefix}.ada_norm.weight": [2048, 4096], + f"{prefix}.ada_norm.bias": [4096], + f"{prefix}.attention.query.weight": [2048, 2048], + f"{prefix}.attention.query.bias": [2048], + f"{prefix}.attention.key.weight": [2048, 2048], + f"{prefix}.attention.key.bias": [2048], + f"{prefix}.attention.value.weight": [2048, 2048], + f"{prefix}.attention.value.bias": [2048], + f"{prefix}.attention.output.weight": [2048, 2048], + f"{prefix}.attention.output.bias": [2048], + f"{prefix}.feed_forward.input.weight": [2048, 8192], + f"{prefix}.feed_forward.input.bias": [8192], + f"{prefix}.feed_forward.output.weight": [8192, 2048], + f"{prefix}.feed_forward.output.bias": [2048], + } + ) + expected.update( + { + "starvla.policy.pi.state.input.weight": [7, 2048], + "starvla.policy.pi.state.input.bias": [2048], + "starvla.policy.pi.state.output.weight": [2048, 2048], + "starvla.policy.pi.state.output.bias": [2048], + "starvla.policy.pi.action.input.weight": [7, 2048], + "starvla.policy.pi.action.input.bias": [2048], + "starvla.policy.pi.action.time_mix.weight": [4096, 2048], + "starvla.policy.pi.action.time_mix.bias": [2048], + "starvla.policy.pi.action.output.weight": [2048, 2048], + "starvla.policy.pi.action.output.bias": [2048], + "starvla.policy.pi.velocity.input.weight": [2048, 2048], + "starvla.policy.pi.velocity.input.bias": [2048], + "starvla.policy.pi.velocity.output.weight": [2048, 7], + "starvla.policy.pi.velocity.output.bias": [7], + "starvla.policy.pi.future_tokens.weight": [2048, 32], + "starvla.policy.pi.action_position.weight": [2048, 1024], + } + ) + if len(expected) != PI_POLICY_TENSOR_COUNT: + raise AssertionError(f"internal legacy PI tensor contract has {len(expected)} tensors") + return expected + + +def expected_pi_v3_policy_tensor_map() -> dict[str, list[int]]: + """Return all PI_v3 GGUF shapes in ggml `ne[]` dimension order.""" + expected = { + "starvla.policy.pi_v3.timestep.input.weight": [256, 1024], + "starvla.policy.pi_v3.timestep.input.bias": [1024], + "starvla.policy.pi_v3.timestep.output.weight": [1024, 1024], + "starvla.policy.pi_v3.timestep.output.bias": [1024], + } + for block in range(PI_V3_BLOCK_COUNT): + prefix = f"starvla.policy.pi_v3.block.{block}" + expected.update( + { + f"{prefix}.ada_norm.weight": [1024, 2048], + f"{prefix}.ada_norm.bias": [2048], + f"{prefix}.attention.query.weight": [1024, 1024], + f"{prefix}.attention.query.bias": [1024], + f"{prefix}.attention.key.weight": [1024, 1024], + f"{prefix}.attention.key.bias": [1024], + f"{prefix}.attention.value.weight": [1024, 1024], + f"{prefix}.attention.value.bias": [1024], + f"{prefix}.attention.output.weight": [1024, 1024], + f"{prefix}.attention.output.bias": [1024], + f"{prefix}.feed_forward.input.weight": [1024, 4096], + f"{prefix}.feed_forward.input.bias": [4096], + f"{prefix}.feed_forward.output.weight": [4096, 1024], + f"{prefix}.feed_forward.output.bias": [1024], + } + ) + expected.update( + { + "starvla.policy.pi_v3.action.input.weight": [7, 1024], + "starvla.policy.pi_v3.action.input.bias": [1024], + "starvla.policy.pi_v3.action.time_mix.weight": [2048, 1024], + "starvla.policy.pi_v3.action.time_mix.bias": [1024], + "starvla.policy.pi_v3.action.output.weight": [1024, 1024], + "starvla.policy.pi_v3.action.output.bias": [1024], + "starvla.policy.pi_v3.velocity.input.weight": [1024, 1024], + "starvla.policy.pi_v3.velocity.input.bias": [1024], + "starvla.policy.pi_v3.velocity.output.weight": [1024, 7], + "starvla.policy.pi_v3.velocity.output.bias": [7], + "starvla.policy.pi_v3.future_tokens.weight": [1024, 32], + "starvla.policy.pi_v3.action_position.weight": [1024, 1024], + } + ) + for projector in range(PI_V3_PROJECTOR_COUNT): + prefix = f"starvla.policy.pi_v3.projector.{projector}" + expected.update( + { + f"{prefix}.norm.weight": [2560], + f"{prefix}.norm.bias": [2560], + f"{prefix}.projection.weight": [2560, 1024], + f"{prefix}.projection.bias": [1024], + } + ) + if len(expected) != PI_V3_POLICY_TENSOR_COUNT: + raise AssertionError(f"internal PI_v3 tensor contract has {len(expected)} tensors") + return expected + + +def metadata_matches(actual: Any, expected: Any) -> bool: + if isinstance(expected, float): + return isinstance(actual, (int, float)) and math.isclose(actual, expected, rel_tol=1e-6, abs_tol=1e-6) + if isinstance(expected, list): + return isinstance(actual, list) and len(actual) == len(expected) and all( + metadata_matches(actual_item, expected_item) + for actual_item, expected_item in zip(actual, expected) + ) + return actual == expected + + +def expect_metadata_field(reader: Any, key: str, expected: Any) -> None: + actual = field_value(reader, key) + if not metadata_matches(actual, expected): + raise StarVLAError(f"GGUF metadata mismatch for {key}: expected {expected!r}, got {actual!r}") + + +def _require_object(value: Any, field: str) -> dict[str, Any]: + if not isinstance(value, dict): + raise StarVLAError(f"pinned Qwen {field} must be a JSON object") + return value + + +def _require_positive_int(value: Any, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise StarVLAError(f"pinned Qwen {field} must be a positive integer") + return value + + +def _special_token_content(value: Any, field: str) -> str | None: + if value is None: + return None + if isinstance(value, str): + return value + if isinstance(value, dict) and isinstance(value.get("content"), str): + return str(value["content"]) + raise StarVLAError(f"pinned Qwen tokenizer_config.json {field} has an unsupported value") + + +def _normalize_merge(merge: Any, index: int) -> str: + if isinstance(merge, str): + return merge + if ( + isinstance(merge, list) + and len(merge) == 2 + and all(isinstance(part, str) for part in merge) + ): + encoded = [ + "".join(chr(ord(character) + 256) if character == " " else character for character in part) + for part in merge + ] + return " ".join(encoded) + raise StarVLAError(f"pinned Qwen tokenizer merge {index} has an unsupported value") + + +def expected_tokenizer_metadata(hf_dir: Path) -> dict[str, Any]: + """Derive llama.cpp's GPT-2 vocabulary metadata from the pinned HF tokenizer files.""" + tokenizer = _load_json(hf_dir / "tokenizer.json") + tokenizer_config = _load_json(hf_dir / "tokenizer_config.json") + config = _load_json(hf_dir / "config.json") + text_config_value = config.get("text_config") + text_config = ( + _require_object(text_config_value, "config.json text_config") + if text_config_value is not None + else config + ) + vocab_size = _require_positive_int(text_config.get("vocab_size"), "text_config.vocab_size") + + model = _require_object(tokenizer.get("model"), "tokenizer.json model") + vocabulary = _require_object(model.get("vocab"), "tokenizer.json model.vocab") + added_tokens = tokenizer.get("added_tokens") + if not isinstance(added_tokens, list): + raise StarVLAError("pinned Qwen tokenizer.json added_tokens must be an array") + + decoder = tokenizer_config.get("added_tokens_decoder") + if not isinstance(decoder, dict): + raise StarVLAError("pinned Qwen tokenizer_config.json added_tokens_decoder must be an object") + decoder_by_id: dict[int, dict[str, Any]] = {} + for raw_id, record in decoder.items(): + if not isinstance(raw_id, str) or not raw_id.isdecimal() or not isinstance(record, dict): + raise StarVLAError("pinned Qwen added_tokens_decoder contains an invalid entry") + token_id = int(raw_id) + if token_id in decoder_by_id: + raise StarVLAError(f"pinned Qwen added_tokens_decoder repeats token id {token_id}") + decoder_by_id[token_id] = record + + tokens = [f"[PAD{token_id}]" for token_id in range(vocab_size)] + token_types = [int(gguf.TokenType.UNUSED)] * vocab_size + assigned_ids: set[int] = set() + token_to_id: dict[str, int] = {} + + def assign(token: Any, token_id: Any, token_type: Any, source: str) -> None: + if not isinstance(token, str): + raise StarVLAError(f"pinned Qwen {source} token must be a string") + if isinstance(token_id, bool) or not isinstance(token_id, int) or not 0 <= token_id < vocab_size: + raise StarVLAError(f"pinned Qwen {source} token id is out of range: {token_id!r}") + if token_id in assigned_ids: + raise StarVLAError(f"pinned Qwen tokenizer repeats token id {token_id}") + if token in token_to_id: + raise StarVLAError(f"pinned Qwen tokenizer repeats token content {token!r}") + assigned_ids.add(token_id) + token_to_id[token] = token_id + tokens[token_id] = token + token_types[token_id] = int(token_type) + + for token, token_id in vocabulary.items(): + assign(token, token_id, gguf.TokenType.NORMAL, "base vocabulary") + + added_by_id: dict[int, dict[str, Any]] = {} + for index, record in enumerate(added_tokens): + if not isinstance(record, dict): + raise StarVLAError(f"pinned Qwen tokenizer added token {index} must be an object") + token_id = record.get("id") + if isinstance(token_id, bool) or not isinstance(token_id, int): + raise StarVLAError(f"pinned Qwen tokenizer added token {index} has an invalid id") + if token_id in added_by_id: + raise StarVLAError(f"pinned Qwen tokenizer repeats added token id {token_id}") + added_by_id[token_id] = record + if set(added_by_id) != set(decoder_by_id): + raise StarVLAError("pinned Qwen tokenizer added_tokens and added_tokens_decoder ids differ") + + for token_id, record in sorted(added_by_id.items()): + decoder_record = decoder_by_id[token_id] + for field in ("content", "normalized", "special"): + if record.get(field) != decoder_record.get(field): + raise StarVLAError( + f"pinned Qwen added token {token_id} disagrees with added_tokens_decoder for {field}" + ) + token = record.get("content") + if not isinstance(token, str): + raise StarVLAError(f"pinned Qwen added token {token_id} has invalid content") + is_control = bool(record.get("special")) or (token.startswith("<|") and token.endswith("|>")) + token_type = gguf.TokenType.CONTROL if is_control else gguf.TokenType.USER_DEFINED + assign(token, token_id, token_type, "added vocabulary") + + raw_merges = model.get("merges") + if not isinstance(raw_merges, list) or not raw_merges: + raise StarVLAError("pinned Qwen tokenizer.json model.merges must be a non-empty array") + merges = [_normalize_merge(merge, index) for index, merge in enumerate(raw_merges)] + + chat_template = tokenizer_config.get("chat_template") + if not isinstance(chat_template, str) or not chat_template: + jinja_path = hf_dir / "chat_template.jinja" + if not jinja_path.is_file(): + raise StarVLAError( + "pinned Qwen tokenizer has no non-empty chat template" + ) + chat_template = jinja_path.read_text(encoding="utf-8") + if not chat_template: + raise StarVLAError("pinned Qwen chat_template.jinja is empty") + chat_template_path = hf_dir / "chat_template.json" + if chat_template_path.is_file(): + template_file = _load_json(chat_template_path).get("chat_template") + if template_file != chat_template: + raise StarVLAError("pinned Qwen chat_template.json disagrees with tokenizer_config.json") + + bos_id = _require_positive_int(text_config.get("bos_token_id"), "text_config.bos_token_id") + eos_id = _require_positive_int(text_config.get("eos_token_id"), "text_config.eos_token_id") + if bos_id >= vocab_size or eos_id >= vocab_size: + raise StarVLAError("pinned Qwen BOS/EOS token id is outside the configured vocabulary") + eos_content = _special_token_content(tokenizer_config.get("eos_token"), "eos_token") + if eos_content is not None and token_to_id.get(eos_content) != eos_id: + raise StarVLAError("pinned Qwen EOS token string and id disagree") + pad_content = _special_token_content(tokenizer_config.get("pad_token"), "pad_token") + if pad_content is None or pad_content not in token_to_id: + raise StarVLAError("pinned Qwen tokenizer has no resolvable padding token") + add_bos = tokenizer_config.get("add_bos_token") + if not isinstance(add_bos, bool): + raise StarVLAError("pinned Qwen tokenizer_config.json add_bos_token must be boolean") + + return { + "tokenizer.ggml.model": "gpt2", + "tokenizer.ggml.pre": "qwen2", + "tokenizer.ggml.tokens": tokens, + "tokenizer.ggml.token_type": token_types, + "tokenizer.ggml.merges": merges, + "tokenizer.ggml.bos_token_id": bos_id, + "tokenizer.ggml.eos_token_id": eos_id, + "tokenizer.ggml.padding_token_id": token_to_id[pad_content], + "tokenizer.ggml.add_bos_token": add_bos, + "tokenizer.chat_template": chat_template, + } + + +def validate_policy_metadata(reader: Any, expected: dict[str, Any]) -> None: + expected_keys = {key for key in expected if key.startswith("starvla.")} + actual_keys = {key for key in reader.fields if key.startswith("starvla.")} + if actual_keys != expected_keys: + raise StarVLAError( + "policy GGUF StarVLA metadata set mismatch; " + f"missing={sorted(expected_keys - actual_keys)}, unexpected={sorted(actual_keys - expected_keys)}" + ) + for key in sorted(expected_keys): + actual = field_value(reader, key) + if not metadata_matches(actual, expected[key]): + raise StarVLAError( + f"policy GGUF metadata mismatch for {key}: expected {expected[key]!r}, got {actual!r}" + ) + expect_field(reader, "general.name", expected["general.name"]) + + +def validate_qwen3vl_image_metadata(reader: Any) -> None: + """Reject fixed-size or incomplete substitutes for the dynamic image contract.""" + expected_keys = set(QWEN3VL_DYNAMIC_IMAGE_METADATA) + actual_keys = {key for key in reader.fields if key.startswith("starvla.image.")} + if actual_keys != expected_keys: + raise StarVLAError( + "Qwen3-VL dynamic image metadata set mismatch; " + f"missing={sorted(expected_keys - actual_keys)}, " + f"unexpected={sorted(actual_keys - expected_keys)}" + ) + for key, expected in sorted(QWEN3VL_DYNAMIC_IMAGE_METADATA.items()): + expect_metadata_field(reader, key, expected) + + +def validate_qwen_vl_image_metadata( + reader: Any, expected_metadata: dict[str, Any], backbone: str +) -> None: + if backbone not in ("qwen3_vl", "qwen2_5_vl"): + raise StarVLAError(f"unsupported Qwen image metadata backbone: {backbone!r}") + expected = { + key: value + for key, value in expected_metadata.items() + if key.startswith("starvla.image.") + } + actual_keys = { + key for key in reader.fields if key.startswith("starvla.image.") + } + if actual_keys != set(expected): + raise StarVLAError( + "Qwen-VL image metadata set mismatch; " + f"missing={sorted(set(expected) - actual_keys)}, " + f"unexpected={sorted(actual_keys - set(expected))}" + ) + for key, value in sorted(expected.items()): + expect_metadata_field(reader, key, value) + + +def validate_dtype_set(reader: Any, requested: str, *, component: str, exact: bool = False) -> dict[str, int]: + requested_type = { + "f32": "F32", + "fp32": "F32", + "f16": "F16", + "bf16": "BF16", + "q8_0": "Q8_0", + }[requested] + counts = Counter(tensor.tensor_type.name for tensor in reader.tensors) + allowed = {requested_type} if exact or requested_type == "F32" else {requested_type, "F32"} + unexpected = set(counts) - allowed + if unexpected or requested_type not in counts: + raise StarVLAError( + f"unexpected {component} GGUF tensor dtypes for {requested}: " + f"counts={dict(sorted(counts.items()))}, allowed={sorted(allowed)}" + ) + return dict(sorted(counts.items())) + + +def _convert_policy_tensor_data(tensor: Any, dtype: str) -> np.ndarray: + array = np.asarray(tensor.detach().float().cpu().numpy(), dtype=np.float32) + if dtype == "fp32": + return array + if dtype == "f16": + return array.astype(np.float16) + if dtype == "bf16": + return gguf.quantize(array, gguf.GGMLQuantizationType.BF16) + raise StarVLAError(f"unsupported policy GGUF dtype: {dtype}") + + +def _first_byte_mismatch(actual: np.ndarray, expected: np.ndarray) -> int | None: + actual_bytes = np.ascontiguousarray(actual).view(np.uint8).reshape(-1) + expected_bytes = np.ascontiguousarray(expected).view(np.uint8).reshape(-1) + if actual_bytes.size != expected_bytes.size: + return min(actual_bytes.size, expected_bytes.size) + chunk_size = 16 * 1024 * 1024 + for offset in range(0, actual_bytes.size, chunk_size): + stop = min(offset + chunk_size, actual_bytes.size) + actual_chunk = actual_bytes[offset:stop] + expected_chunk = expected_bytes[offset:stop] + if not np.array_equal(actual_chunk, expected_chunk): + mismatch = np.flatnonzero(actual_chunk != expected_chunk) + return offset + int(mismatch[0]) + return None + + +def validate_policy_tensor_bytes( + tensors: dict[str, Any], + policy_dir: Path, + dtype: str, + tensor_name_map: dict[str, str] | None = None, + component_label: str = "OFT", +) -> None: + tensor_name_map = OFT_TENSOR_MAP if tensor_name_map is None else tensor_name_map + source_tensors = load_policy_tensors(policy_dir) + missing = sorted(set(tensor_name_map) - set(source_tensors)) + if missing: + raise StarVLAError( + f"staged {component_label} policy is missing runtime tensors: {missing}" + ) + if set(tensors) != set(tensor_name_map.values()): + raise StarVLAError( + f"{component_label} policy GGUF tensor names do not match the canonical " + f"{len(tensor_name_map)}-tensor map" + ) + for source_name, destination_name in tensor_name_map.items(): + expected = _convert_policy_tensor_data(source_tensors[source_name], dtype) + actual = np.asarray(tensors[destination_name].data) + if actual.nbytes != expected.nbytes: + raise StarVLAError( + f"policy GGUF tensor byte size mismatch for {destination_name}: " + f"expected {expected.nbytes}, got {actual.nbytes}" + ) + mismatch = _first_byte_mismatch(actual, expected) + if mismatch is not None: + raise StarVLAError( + f"policy GGUF tensor content mismatch for {destination_name} at byte offset {mismatch}" + ) + del expected + del source_tensors + + +def validate_text( + reader: Any, + bundle_uuid: str, + dtype: str, + hf_dir: Path, + *, + require_oft_action_token: bool = True, + backbone: str = "qwen3_vl", +) -> dict[str, Any]: + if backbone == "qwen3_vl": + architecture = "qwen3vl" + expect_field(reader, "general.architecture", architecture) + expect_field(reader, "qwen3vl.context_length", 262144) + expect_field(reader, "qwen3vl.embedding_length", 2560) + expect_field(reader, "qwen3vl.feed_forward_length", 9728) + expect_field(reader, "qwen3vl.block_count", 36) + expect_field(reader, "qwen3vl.attention.head_count", 32) + expect_field(reader, "qwen3vl.attention.head_count_kv", 8) + expect_field(reader, "qwen3vl.attention.key_length", 128) + expect_field(reader, "qwen3vl.attention.value_length", 128) + expect_metadata_field( + reader, "qwen3vl.attention.layer_norm_rms_epsilon", 1e-6 + ) + expect_field(reader, "qwen3vl.rope.dimension_sections", [24, 20, 20, 0]) + expect_metadata_field(reader, "qwen3vl.rope.freq_base", 5_000_000.0) + expect_field(reader, "qwen3vl.n_deepstack_layers", 3) + vocab_size = 151936 + elif backbone == "qwen2_5_vl": + architecture = "qwen2vl" + config = _load_json(hf_dir / "config.json") + text_config_value = config.get("text_config") + text_config = ( + _require_object(text_config_value, "config.json text_config") + if text_config_value is not None + else config + ) + vocab_size = _require_positive_int( + text_config.get("vocab_size"), "text_config.vocab_size" + ) + expect_field(reader, "general.architecture", architecture) + expect_field(reader, "qwen2vl.context_length", 128000) + expect_field(reader, "qwen2vl.embedding_length", 2048) + expect_field(reader, "qwen2vl.feed_forward_length", 11008) + expect_field(reader, "qwen2vl.block_count", 36) + expect_field(reader, "qwen2vl.attention.head_count", 16) + expect_field(reader, "qwen2vl.attention.head_count_kv", 2) + expect_metadata_field( + reader, "qwen2vl.attention.layer_norm_rms_epsilon", 1e-6 + ) + expect_field(reader, "qwen2vl.rope.dimension_sections", [16, 24, 24, 0]) + expect_metadata_field(reader, "qwen2vl.rope.freq_base", 1_000_000.0) + if "qwen2vl.n_deepstack_layers" in reader.fields: + raise StarVLAError("Qwen2.5-VL text GGUF unexpectedly enables DeepStack") + else: + raise StarVLAError(f"unsupported Qwen text backbone: {backbone!r}") + tokenizer_metadata = expected_tokenizer_metadata(hf_dir) + for key, expected in tokenizer_metadata.items(): + if isinstance(expected, list): + expect_sequence_field(reader, key, expected) + else: + expect_field(reader, key, expected) + if require_oft_action_token: + action_token_id = resolve_action_token_id(hf_dir) + if action_token_id != OFT_ACTION_TOKEN_ID: + raise StarVLAError( + f"Qwen action token id mismatch: expected {OFT_ACTION_TOKEN_ID}, got {action_token_id}" + ) + tensors = tensor_map(reader) + expect_complete_tensor_map( + tensors, + expected_text_tensor_map(backbone, vocab_size), + "Qwen text", + ) + return { + "architecture": architecture, + "tensor_count": len(tensors), + "dtypes": validate_dtype_set(reader, dtype, component="text"), + } + + +def validate_mmproj( + reader: Any, + bundle_uuid: str, + dtype: str, + hf_dir: Path, + *, + backbone: str = "qwen3_vl", +) -> dict[str, Any]: + config = _load_json(hf_dir / "config.json") + text_config_value = config.get("text_config") + text_config = ( + _require_object(text_config_value, "config.json text_config") + if text_config_value is not None + else config + ) + vision_config = _require_object(config.get("vision_config"), "config.json vision_config") + preprocessor = _load_json(hf_dir / "preprocessor_config.json") + patch_size = _require_positive_int(vision_config.get("patch_size"), "vision_config.patch_size") + image_mean = preprocessor.get("image_mean") + image_std = preprocessor.get("image_std") + if not isinstance(image_mean, list) or not isinstance(image_std, list): + raise StarVLAError("pinned Qwen preprocessor image_mean/image_std must be arrays") + + expect_field(reader, "general.architecture", "clip") + expect_field(reader, "general.source.uuid", bundle_uuid) + expect_field(reader, "clip.has_vision_encoder", True) + expect_field(reader, "clip.vision.patch_size", patch_size) + expect_field(reader, "clip.vision.embedding_length", vision_config.get("hidden_size")) + expect_field(reader, "clip.vision.feed_forward_length", vision_config.get("intermediate_size")) + expect_field(reader, "clip.vision.projection_dim", text_config.get("hidden_size")) + expect_field(reader, "clip.vision.block_count", vision_config.get("depth")) + expect_field(reader, "clip.vision.attention.head_count", vision_config.get("num_heads")) + expect_metadata_field(reader, "clip.vision.attention.layer_norm_epsilon", text_config.get("rms_norm_eps")) + expect_metadata_field(reader, "clip.vision.image_mean", image_mean) + expect_metadata_field(reader, "clip.vision.image_std", image_std) + if backbone == "qwen3_vl": + num_positions = _require_positive_int( + vision_config.get("num_position_embeddings"), + "vision_config.num_position_embeddings", + ) + positions_per_side = math.isqrt(num_positions) + if positions_per_side * positions_per_side != num_positions: + raise StarVLAError("pinned Qwen vision position count is not square") + deepstack_indices = vision_config.get("deepstack_visual_indexes") + if not isinstance(deepstack_indices, list) or any( + isinstance(index, bool) or not isinstance(index, int) + for index in deepstack_indices + ): + raise StarVLAError( + "pinned Qwen vision_config.deepstack_visual_indexes must be an integer array" + ) + expect_field(reader, "clip.projector_type", "qwen3vl_merger") + expect_field( + reader, "clip.vision.image_size", positions_per_side * patch_size + ) + expect_field(reader, "clip.use_gelu", True) + expect_field( + reader, + "clip.vision.spatial_merge_size", + vision_config.get("spatial_merge_size"), + ) + deepstack = field_value(reader, "clip.vision.is_deepstack_layers") + if ( + len(deepstack) != int(vision_config.get("depth", -1)) + or [ + index for index, enabled in enumerate(deepstack) if enabled + ] + != deepstack_indices + ): + raise StarVLAError( + f"Qwen vision DeepStack layer mismatch: {deepstack}" + ) + elif backbone == "qwen2_5_vl": + expect_field(reader, "clip.projector_type", "qwen2.5vl_merger") + expect_field(reader, "clip.vision.image_size", 560) + expect_field(reader, "clip.use_silu", True) + expect_field(reader, "clip.vision.n_wa_pattern", 8) + if "clip.vision.is_deepstack_layers" in reader.fields: + raise StarVLAError( + "Qwen2.5-VL mmproj unexpectedly contains DeepStack metadata" + ) + else: + raise StarVLAError(f"unsupported Qwen mmproj backbone: {backbone!r}") + tensors = tensor_map(reader) + expect_complete_tensor_map( + tensors, expected_mmproj_tensor_map(backbone), "Qwen mmproj" + ) + return { + "architecture": "clip", + "tensor_count": len(tensors), + "dtypes": validate_dtype_set(reader, dtype, component="mmproj"), + } + + +def validate_policy( + reader: Any, + bundle_uuid: str, + dtype: str, + policy_dir: Path, + text_filename: str, + mmproj_filename: str, + expected_metadata: dict[str, Any], + framework: str = "oft", + backbone: str = "qwen3_vl", +) -> dict[str, Any]: + if framework not in ("oft", "groot", "pi", "pi_v3"): + raise StarVLAError(f"unsupported policy framework for validation: {framework}") + model_type = str(expected_metadata["starvla.model_type"]) + expect_field(reader, "general.architecture", "starvla-policy") + expect_field(reader, "general.source.uuid", bundle_uuid) + expect_field(reader, "starvla.bundle.uuid", bundle_uuid) + expect_field(reader, "starvla.framework", framework) + expect_field(reader, "starvla.model_type", model_type) + expect_field(reader, "starvla.component.text.filename", text_filename) + expect_field(reader, "starvla.component.mmproj.filename", mmproj_filename) + expect_field(reader, "starvla.backbone.arch", backbone) + if framework == "oft": + expect_field(reader, "starvla.prompt.action_token_id", 146663) + elif framework == "groot": + expect_field(reader, "starvla.groot.timestep_ids", [0, 250, 500, 750]) + elif framework == "pi": + expect_field( + reader, + "starvla.conditioning.hidden_tuple_indices", + list(range(21, 37)), + ) + expect_field(reader, "starvla.pi.timestep_ids", [0, 250, 500, 750]) + expect_field( + reader, "starvla.image.framework_inference_pre_resize_width", 224 + ) + expect_field( + reader, "starvla.image.framework_inference_pre_resize_height", 224 + ) + validate_qwen_vl_image_metadata(reader, expected_metadata, backbone) + expect_field(reader, "starvla.action.dimension", 7) + expect_field(reader, "starvla.action.horizon", 16) + expect_field(reader, "starvla.normalization.profile_count", 2) + expect_field( + reader, + "starvla.normalization.profile_keys", + expected_metadata["starvla.normalization.profile_keys"], + ) + validate_policy_metadata(reader, expected_metadata) + tensors = tensor_map(reader) + if framework == "oft": + tensor_name_map = OFT_TENSOR_MAP + input_dim = 2048 if backbone == "qwen2_5_vl" else 2560 + hidden_dim = 4096 if backbone == "qwen2_5_vl" else 5120 + expected_shapes = { + "starvla.policy.oft.input_norm.weight": [input_dim], + "starvla.policy.oft.input_norm.bias": [input_dim], + "starvla.policy.oft.input_proj.weight": [input_dim, hidden_dim], + "starvla.policy.oft.input_proj.bias": [hidden_dim], + "starvla.policy.oft.output_norm.weight": [hidden_dim], + "starvla.policy.oft.output_norm.bias": [hidden_dim], + "starvla.policy.oft.output_proj.weight": [hidden_dim, 7], + "starvla.policy.oft.output_proj.bias": [7], + } + for block in (0, 1): + expected_shapes[f"starvla.policy.oft.block.{block}.norm.weight"] = [hidden_dim] + expected_shapes[f"starvla.policy.oft.block.{block}.norm.bias"] = [hidden_dim] + expected_shapes[f"starvla.policy.oft.block.{block}.linear.weight"] = [ + hidden_dim, + hidden_dim, + ] + expected_shapes[f"starvla.policy.oft.block.{block}.linear.bias"] = [hidden_dim] + elif framework == "groot": + tensor_name_map = GROOT_TENSOR_MAP + expected_shapes = expected_groot_policy_tensor_map( + 2048 if backbone == "qwen2_5_vl" else 2560 + ) + elif framework == "pi": + tensor_name_map = PI_TENSOR_MAP + expected_shapes = expected_pi_policy_tensor_map() + else: + tensor_name_map = PI_V3_TENSOR_MAP + expected_shapes = expected_pi_v3_policy_tensor_map() + expect_complete_tensor_map(tensors, expected_shapes, f"{framework.upper()} policy") + dtype_counts = validate_dtype_set(reader, dtype, component="policy", exact=True) + validate_policy_tensor_bytes( + tensors, + policy_dir, + dtype, + tensor_name_map=tensor_name_map, + component_label=framework.upper(), + ) + return { + "architecture": "starvla-policy", + "tensor_count": len(tensors), + "dtypes": dtype_counts, + } + + +def component_record(path: Path, validation: dict[str, Any]) -> dict[str, Any]: + return { + "filename": path.name, + "size": path.stat().st_size, + "sha256": sha256_file(path), + **validation, + } + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--variant", + default="oft", + choices=( + "oft", + "groot", + "pi_v3", + "qwen25_oft", + "qwen25_groot", + "qwen25_pi", + ), + ) + parser.add_argument("--text", type=Path, required=True) + parser.add_argument("--mmproj", type=Path, required=True) + parser.add_argument("--policy", type=Path, required=True) + parser.add_argument("--hf-dir", type=Path, required=True) + parser.add_argument("--policy-dir", type=Path, required=True) + parser.add_argument("--surgery-manifest", type=Path, required=True) + parser.add_argument("--catalog", type=Path, default=DEFAULT_CATALOG) + parser.add_argument( + "--text-dtype", + choices=("f32", "f16", "bf16", "q8_0"), + default=DEFAULT_TEXT_DTYPE, + ) + parser.add_argument( + "--mmproj-dtype", + choices=("f32", "f16", "bf16", "q8_0"), + default=DEFAULT_MMPROJ_DTYPE, + ) + parser.add_argument( + "--policy-dtype", + choices=("fp32", "f16", "bf16"), + default=DEFAULT_POLICY_DTYPE, + ) + parser.add_argument("--output", type=Path, required=True) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + try: + if args.output.exists() or args.output.is_symlink(): + raise StarVLAError(f"refusing to overwrite existing output: {args.output}") + for path in (args.text, args.mmproj, args.policy): + if not path.is_file() or path.stat().st_size == 0: + raise StarVLAError(f"missing or empty bundle component: {path}") + surgery_manifest = _load_json(args.surgery_manifest) + catalog = load_catalog(args.catalog) + variant = get_variant(catalog, args.variant) + framework = str(variant["framework"]) + backbone = str(variant.get("backbone", "qwen3_vl")) + validate_official_surgery_manifest(surgery_manifest, variant, catalog) + verify_staged_assets(args.hf_dir, surgery_manifest.get("qwen_assets", {}), component="Qwen") + verify_staged_assets(args.policy_dir, surgery_manifest.get("policy_assets", {}), component="policy") + verify_staged_components_against_checkpoint( + { + "vlm": (args.hf_dir, surgery_manifest.get("vlm_output", {})), + "policy": (args.policy_dir, surgery_manifest.get("policy_output", {})), + }, + surgery_manifest, + variant, + ) + bundle_uuid = str(surgery_manifest["bundle_uuid"]) + if framework == "oft": + oft_dimensions = ( + {"input_dim": 2048, "hidden_dim": 4096, "action_dim": 7} + if backbone == "qwen2_5_vl" + else {"input_dim": 2560, "hidden_dim": 5120, "action_dim": 7} + ) + expected_policy_metadata = build_oft_metadata( + args.policy_dir, + args.hf_dir, + variant, + surgery_manifest, + oft_dimensions, + OFT_ACTION_TOKEN_ID, + args.text.name, + args.mmproj.name, + ) + elif framework == "groot": + groot_dimensions = dict( + GROOT_OFFICIAL_DIMENSIONS_BY_BACKBONE[backbone] + ) + expected_policy_metadata = build_groot_metadata( + args.policy_dir, + args.hf_dir, + variant, + surgery_manifest, + groot_dimensions, + args.text.name, + args.mmproj.name, + ) + elif framework == "pi": + expected_policy_metadata = build_pi_metadata( + args.policy_dir, + args.hf_dir, + variant, + surgery_manifest, + dict(PI_OFFICIAL_DIMENSIONS), + args.text.name, + args.mmproj.name, + ) + else: + expected_policy_metadata = build_pi_v3_metadata( + args.policy_dir, + args.hf_dir, + variant, + surgery_manifest, + dict(PI_V3_OFFICIAL_DIMENSIONS), + args.text.name, + args.mmproj.name, + ) + + text_reader = gguf.GGUFReader(args.text) + mmproj_reader = gguf.GGUFReader(args.mmproj) + policy_reader = gguf.GGUFReader(args.policy) + text_validation = validate_text( + text_reader, + bundle_uuid, + args.text_dtype, + args.hf_dir, + require_oft_action_token=framework == "oft", + backbone=backbone, + ) + mmproj_validation = validate_mmproj( + mmproj_reader, + bundle_uuid, + args.mmproj_dtype, + args.hf_dir, + backbone=backbone, + ) + policy_validation = validate_policy( + policy_reader, + bundle_uuid, + args.policy_dtype, + args.policy_dir, + args.text.name, + args.mmproj.name, + expected_policy_metadata, + framework=framework, + backbone=backbone, + ) + del text_reader, mmproj_reader, policy_reader + + source_tensors = surgery_manifest["tensors"] + role_counts = Counter(record["role"] for record in source_tensors) + expected_role_counts = Counter( + { + "text": int(variant["expected"]["text_tensors"]), + "visual": int(variant["expected"]["visual_tensors"]), + "policy": int(variant["expected"]["policy_tensors"]), + "lm_head": int(variant["expected"]["lm_head_tensors"]), + } + ) + if role_counts != expected_role_counts: + raise StarVLAError(f"unexpected surgery source tensor coverage: {dict(role_counts)}") + manifest = { + "schema_version": 1, + "variant": args.variant, + "model_type": variant["model_type"], + "bundle_uuid": bundle_uuid, + "source": surgery_manifest["source"], + "source_tensor_roles": dict(sorted(role_counts.items())), + "surgery_manifest": { + "filename": args.surgery_manifest.name, + "size": args.surgery_manifest.stat().st_size, + "sha256": sha256_file(args.surgery_manifest), + }, + "components": { + "text": component_record(args.text, text_validation), + "mmproj": component_record(args.mmproj, mmproj_validation), + "policy": component_record(args.policy, policy_validation), + }, + } + atomic_write_json(args.output, manifest, overwrite=False) + print(f"conversion manifest: {args.output}") + return 0 + except (StarVLAError, OSError, ValueError, TypeError, KeyError) as exc: + print(f"error: {exc}", file=sys.stderr) + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/llama_cpp/apply_starvla_patches.sh b/tools/llama_cpp/apply_starvla_patches.sh new file mode 100755 index 0000000..be471c1 --- /dev/null +++ b/tools/llama_cpp/apply_starvla_patches.sh @@ -0,0 +1,115 @@ +#!/usr/bin/env bash +set -euo pipefail + +readonly EXPECTED_LLAMA_COMMIT="3e941b813b1acbbf06c2203a94ceb33d84748c1e" +readonly SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +readonly REPO_ROOT="$(git -C "${SCRIPT_DIR}" rev-parse --show-toplevel)" +readonly LLAMA_DIR="${LLAMA_CPP_DIR:-${REPO_ROOT}/third_party/llama.cpp}" +readonly PATCH_DIR="${REPO_ROOT}/patches/llama.cpp" +readonly PATCHES=( + "${PATCH_DIR}/0001-qwen3vl-vision-parity.patch" + "${PATCH_DIR}/0002-per-context-native-graph-control.patch" +) + +usage() { + cat <<'EOF' +Usage: tools/llama_cpp/apply_starvla_patches.sh [--check|--revert] + +With no option, apply the StarVLA patches to third_party/llama.cpp. + --check Validate the pinned revision and report patch state. + --revert Remove an already applied complete patch set. + +Set LLAMA_CPP_DIR to validate or patch another checkout of the pinned revision. +EOF +} + +mode="apply" +case "${1:-}" in + "") ;; + --check) mode="check" ;; + --revert) mode="revert" ;; + -h|--help) usage; exit 0 ;; + *) usage >&2; exit 2 ;; +esac + +if [[ ! -d "${LLAMA_DIR}/.git" && ! -f "${LLAMA_DIR}/.git" ]]; then + echo "error: llama.cpp checkout not found: ${LLAMA_DIR}" >&2 + exit 1 +fi + +actual_commit="$(git -C "${LLAMA_DIR}" rev-parse HEAD)" +if [[ "${actual_commit}" != "${EXPECTED_LLAMA_COMMIT}" ]]; then + echo "error: unsupported llama.cpp revision" >&2 + echo " expected: ${EXPECTED_LLAMA_COMMIT}" >&2 + echo " actual: ${actual_commit}" >&2 + exit 1 +fi + +states=() +for patch in "${PATCHES[@]}"; do + if git -C "${LLAMA_DIR}" apply --reverse --check "${patch}" >/dev/null 2>&1; then + states+=("applied") + elif git -C "${LLAMA_DIR}" apply --check "${patch}" >/dev/null 2>&1; then + states+=("pending") + else + echo "error: patch is neither cleanly applicable nor already applied: ${patch}" >&2 + exit 1 + fi +done + +all_pending=true +all_applied=true +for state in "${states[@]}"; do + [[ "${state}" == "pending" ]] || all_pending=false + [[ "${state}" == "applied" ]] || all_applied=false +done + +if [[ "${mode}" == "check" ]]; then + for i in "${!PATCHES[@]}"; do + printf '%-8s %s\n' "${states[$i]}" "${PATCHES[$i]#${REPO_ROOT}/}" + done + if ! ${all_pending} && ! ${all_applied}; then + echo "error: partial patch set detected" >&2 + exit 1 + fi + exit 0 +fi + +if [[ "${mode}" == "apply" ]]; then + if ${all_applied}; then + echo "StarVLA llama.cpp patches are already applied." + exit 0 + fi + if ! ${all_pending}; then + echo "error: refusing to apply a partial patch set" >&2 + exit 1 + fi + if [[ -n "$(git -C "${LLAMA_DIR}" status --porcelain)" ]]; then + echo "error: refusing to patch a dirty llama.cpp checkout" >&2 + exit 1 + fi + for patch in "${PATCHES[@]}"; do + git -C "${LLAMA_DIR}" apply --check "${patch}" + done + for patch in "${PATCHES[@]}"; do + git -C "${LLAMA_DIR}" apply "${patch}" + echo "applied ${patch#${REPO_ROOT}/}" + done + exit 0 +fi + +if ${all_pending}; then + echo "StarVLA llama.cpp patches are not applied." + exit 0 +fi +if ! ${all_applied}; then + echo "error: refusing to revert a partial patch set" >&2 + exit 1 +fi +for ((i=${#PATCHES[@]} - 1; i >= 0; --i)); do + git -C "${LLAMA_DIR}" apply --reverse --check "${PATCHES[$i]}" +done +for ((i=${#PATCHES[@]} - 1; i >= 0; --i)); do + git -C "${LLAMA_DIR}" apply --reverse "${PATCHES[$i]}" + echo "reverted ${PATCHES[$i]#${REPO_ROOT}/}" +done From 2ea489516b0dd92acea955610eac6bc77a6943e1 Mon Sep 17 00:00:00 2001 From: JJJYmmm <1650675829@qq.com> Date: Mon, 10 Aug 2026 12:42:56 +0800 Subject: [PATCH 02/11] starvla: add Qwen3-VL OFT policy --- src/models/starvla/oft_image_preprocess.cpp | 363 ++++++ src/models/starvla/oft_image_preprocess.h | 34 + src/models/starvla/oft_policy.cpp | 539 ++++++++ src/models/starvla/oft_policy.h | 65 + src/models/starvla/oft_prompt.cpp | 165 +++ src/models/starvla/oft_prompt.h | 35 + tests/starvla/oft_image_preprocess_test.cpp | 162 +++ tests/starvla/oft_prompt_test.cpp | 81 ++ .../starvla/generate_starvla_oft_golden.py | 928 +++++++++++++ .../starvla/serve_starvla_oft_reference.py | 1152 +++++++++++++++++ 10 files changed, 3524 insertions(+) create mode 100644 src/models/starvla/oft_image_preprocess.cpp create mode 100644 src/models/starvla/oft_image_preprocess.h create mode 100644 src/models/starvla/oft_policy.cpp create mode 100644 src/models/starvla/oft_policy.h create mode 100644 src/models/starvla/oft_prompt.cpp create mode 100644 src/models/starvla/oft_prompt.h create mode 100644 tests/starvla/oft_image_preprocess_test.cpp create mode 100644 tests/starvla/oft_prompt_test.cpp create mode 100644 tools/hf2gguf/starvla/generate_starvla_oft_golden.py create mode 100644 tools/hf2gguf/starvla/serve_starvla_oft_reference.py diff --git a/src/models/starvla/oft_image_preprocess.cpp b/src/models/starvla/oft_image_preprocess.cpp new file mode 100644 index 0000000..3192fcd --- /dev/null +++ b/src/models/starvla/oft_image_preprocess.cpp @@ -0,0 +1,363 @@ +#include "models/starvla/oft_image_preprocess.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace robotcpp::starvla { +namespace { + +struct RGBImage { + int width = 0; + int height = 0; + std::vector pixels; +}; + +struct FilterTable { + int kernel_size = 0; + int precision = 0; + std::vector first; + std::vector count; + std::vector weights; +}; + +double keys_cubic(double value) { + constexpr double a = -0.5; + value = std::fabs(value); + if (value < 1.0) { + return ((a + 2.0) * value - (a + 3.0)) * value * value + 1.0; + } + if (value < 2.0) { + return (((value - 5.0) * value + 8.0) * value - 4.0) * a; + } + return 0.0; +} + +bool validate_resize(const uint8_t * source, int source_width, int source_height, + int source_stride, int target_width, int target_height, + std::string & error) { + error.clear(); + if (source == nullptr || source_width <= 0 || source_height <= 0 || target_width <= 0 || + target_height <= 0) { + error = "StarVLA image resize received an invalid image or dimension"; + return false; + } + const int packed_stride = source_width * 3; + if (source_stride != 0 && source_stride < packed_stride) { + error = "StarVLA image stride is smaller than a packed RGB row"; + return false; + } + const uint64_t output_bytes = static_cast(target_width) * + static_cast(target_height) * 3; + if (output_bytes > static_cast(std::numeric_limits::max())) { + error = "StarVLA resized image is too large"; + return false; + } + return true; +} + +RGBImage pack_source(const uint8_t * source, int width, int height, int stride) { + RGBImage image; + image.width = width; + image.height = height; + const size_t row_bytes = static_cast(width) * 3; + const size_t actual_stride = stride > 0 ? static_cast(stride) : row_bytes; + image.pixels.resize(row_bytes * static_cast(height)); + for (int row = 0; row < height; ++row) { + std::copy(source + static_cast(row) * actual_stride, + source + static_cast(row) * actual_stride + row_bytes, + image.pixels.begin() + static_cast(static_cast(row) * row_bytes)); + } + return image; +} + +FilterTable make_filter_table(int input_size, int output_size, bool pillow_precision) { + const double scale = static_cast(input_size) / static_cast(output_size); + const double filter_scale = std::max(scale, 1.0); + const double support = 2.0 * filter_scale; + + FilterTable table; + table.kernel_size = static_cast(std::ceil(support)) * 2 + 1; + table.first.resize(static_cast(output_size)); + table.count.resize(static_cast(output_size)); + std::vector floating_weights(static_cast(output_size) * + static_cast(table.kernel_size), 0.0); + double maximum_weight = 0.0; + + for (int output = 0; output < output_size; ++output) { + const double center = (static_cast(output) + 0.5) * scale; + const int first = std::max(static_cast(center - support + 0.5), 0); + const int end = std::min(static_cast(center + support + 0.5), input_size); + const int count = std::max(0, std::min(end - first, table.kernel_size)); + table.first[static_cast(output)] = first; + table.count[static_cast(output)] = count; + + double sum = 0.0; + for (int index = 0; index < count; ++index) { + const double distance = + (static_cast(index + first) - center + 0.5) / filter_scale; + const double weight = keys_cubic(distance); + floating_weights[static_cast(output) * table.kernel_size + index] = weight; + sum += weight; + } + if (sum != 0.0) { + for (int index = 0; index < count; ++index) { + double & weight = + floating_weights[static_cast(output) * table.kernel_size + index]; + weight /= sum; + maximum_weight = std::max(maximum_weight, weight); + } + } + } + + if (pillow_precision) { + table.precision = 22; + } else { + for (table.precision = 0; table.precision < 22; ++table.precision) { + const int next = static_cast( + 0.5 + maximum_weight * static_cast(uint32_t{1} << (table.precision + 1))); + if (next >= (1 << 15)) { + break; + } + } + } + + const double multiplier = static_cast(uint32_t{1} << table.precision); + table.weights.resize(floating_weights.size()); + for (size_t index = 0; index < floating_weights.size(); ++index) { + const double scaled = floating_weights[index] * multiplier; + table.weights[index] = static_cast(scaled < 0.0 ? scaled - 0.5 : scaled + 0.5); + } + return table; +} + +uint8_t fixed_point_pixel(int64_t accumulator, int precision) { + const int64_t value = accumulator >> precision; + return static_cast(std::max(0, std::min(255, value))); +} + +RGBImage resize_horizontal(const RGBImage & source, int target_width, + const FilterTable & table) { + RGBImage target; + target.width = target_width; + target.height = source.height; + target.pixels.resize(static_cast(target.width) * target.height * 3); + const int64_t rounding = int64_t{1} << (table.precision - 1); + for (int row = 0; row < source.height; ++row) { + for (int column = 0; column < target.width; ++column) { + const int first = table.first[static_cast(column)]; + const int count = table.count[static_cast(column)]; + for (int channel = 0; channel < 3; ++channel) { + int64_t accumulator = rounding; + for (int index = 0; index < count; ++index) { + const size_t source_index = + (static_cast(row) * source.width + first + index) * 3 + channel; + const int32_t weight = + table.weights[static_cast(column) * table.kernel_size + index]; + accumulator += static_cast(source.pixels[source_index]) * weight; + } + const size_t target_index = + (static_cast(row) * target.width + column) * 3 + channel; + target.pixels[target_index] = fixed_point_pixel(accumulator, table.precision); + } + } + } + return target; +} + +RGBImage resize_vertical(const RGBImage & source, int target_height, + const FilterTable & table) { + RGBImage target; + target.width = source.width; + target.height = target_height; + target.pixels.resize(static_cast(target.width) * target.height * 3); + const int64_t rounding = int64_t{1} << (table.precision - 1); + for (int row = 0; row < target.height; ++row) { + const int first = table.first[static_cast(row)]; + const int count = table.count[static_cast(row)]; + for (int column = 0; column < target.width; ++column) { + for (int channel = 0; channel < 3; ++channel) { + int64_t accumulator = rounding; + for (int index = 0; index < count; ++index) { + const size_t source_index = + (static_cast(first + index) * source.width + column) * 3 + channel; + const int32_t weight = + table.weights[static_cast(row) * table.kernel_size + index]; + accumulator += static_cast(source.pixels[source_index]) * weight; + } + const size_t target_index = + (static_cast(row) * target.width + column) * 3 + channel; + target.pixels[target_index] = fixed_point_pixel(accumulator, table.precision); + } + } + } + return target; +} + +bool resize_rgb(const uint8_t * source, int source_width, int source_height, int source_stride, + int target_width, int target_height, bool pillow_precision, + std::vector & target, std::string & error) { + target.clear(); + if (!validate_resize(source, source_width, source_height, source_stride, target_width, + target_height, error)) { + return false; + } + + RGBImage current = pack_source(source, source_width, source_height, source_stride); + if (source_width != target_width) { + current = resize_horizontal(current, target_width, + make_filter_table(source_width, target_width, pillow_precision)); + } + if (source_height != target_height) { + current = resize_vertical(current, target_height, + make_filter_table(source_height, target_height, pillow_precision)); + } + target = std::move(current.pixels); + return true; +} + +} // namespace + +bool resize_pillow_bicubic_rgb(const uint8_t * source, int source_width, int source_height, + int source_stride, int target_width, int target_height, + std::vector & target, std::string & error) { + return resize_rgb(source, source_width, source_height, source_stride, target_width, + target_height, true, target, error); +} + +bool resize_torchvision_bicubic_aa_rgb(const uint8_t * source, int source_width, + int source_height, int source_stride, int target_width, + int target_height, std::vector & target, + std::string & error) { + return resize_rgb(source, source_width, source_height, source_stride, target_width, + target_height, false, target, error); +} + +bool qwen3vl_smart_resize_dimensions(int source_width, int source_height, int factor, + int min_pixels, int max_pixels, int & target_width, + int & target_height, std::string & error) { + target_width = 0; + target_height = 0; + error.clear(); + if (source_width <= 0 || source_height <= 0 || factor <= 0 || min_pixels <= 0 || + max_pixels < min_pixels) { + error = "Qwen3-VL smart resize received an invalid dimension or pixel bound"; + return false; + } + const int minimum_side = std::min(source_width, source_height); + const int maximum_side = std::max(source_width, source_height); + if (static_cast(maximum_side) / minimum_side > 200.0) { + error = "Qwen3-VL smart resize requires an absolute aspect ratio of at most 200"; + return false; + } + + // Python round() uses ties-to-even. Integer quotient/remainder arithmetic + // makes the common first smart_resize step independent of the host FP mode. + const auto round_div_ties_to_even = [](int value, int divisor) -> int64_t { + const int64_t quotient = value / divisor; + const int64_t remainder = value % divisor; + const int64_t doubled = remainder * 2; + if (doubled < divisor || (doubled == divisor && quotient % 2 == 0)) { + return quotient; + } + return quotient + 1; + }; + + int64_t resized_height = round_div_ties_to_even(source_height, factor) * factor; + int64_t resized_width = round_div_ties_to_even(source_width, factor) * factor; + const int64_t source_pixels = static_cast(source_height) * source_width; + const int64_t rounded_pixels = resized_height * resized_width; + if (rounded_pixels > max_pixels) { + const double beta = std::sqrt(static_cast(source_pixels) / max_pixels); + resized_height = std::max( + factor, static_cast(std::floor(source_height / beta / factor)) * factor); + resized_width = std::max( + factor, static_cast(std::floor(source_width / beta / factor)) * factor); + } else if (rounded_pixels < min_pixels) { + const double beta = std::sqrt(static_cast(min_pixels) / source_pixels); + resized_height = + static_cast(std::ceil(source_height * beta / factor)) * factor; + resized_width = + static_cast(std::ceil(source_width * beta / factor)) * factor; + } + + if (resized_width <= 0 || resized_height <= 0 || + resized_width > std::numeric_limits::max() || + resized_height > std::numeric_limits::max() || + resized_width > std::numeric_limits::max() / resized_height) { + error = "Qwen3-VL smart resize produced an unsupported output dimension"; + return false; + } + target_width = static_cast(resized_width); + target_height = static_cast(resized_height); + return true; +} + +bool preprocess_qwen3vl_rgb(const uint8_t * source, int source_width, int source_height, + int channels, int source_stride, int patch_size, + int spatial_merge_size, int min_pixels, int max_pixels, + std::vector & target, int & target_width, + int & target_height, int & image_token_count, + std::string & error) { + target.clear(); + target_width = 0; + target_height = 0; + image_token_count = 0; + if (channels != 3) { + error = "Qwen3-VL input image must be RGB"; + return false; + } + if (patch_size <= 0 || spatial_merge_size <= 0 || + patch_size > std::numeric_limits::max() / spatial_merge_size) { + error = "Qwen3-VL patch or spatial merge size is invalid"; + return false; + } + const int factor = patch_size * spatial_merge_size; + if (!qwen3vl_smart_resize_dimensions(source_width, source_height, factor, min_pixels, + max_pixels, target_width, target_height, error)) { + return false; + } + if (!resize_torchvision_bicubic_aa_rgb(source, source_width, source_height, source_stride, + target_width, target_height, target, error)) { + target_width = 0; + target_height = 0; + return false; + } + const int64_t grid_width = target_width / factor; + const int64_t grid_height = target_height / factor; + const int64_t tokens = grid_width * grid_height; + if (tokens <= 0 || tokens > std::numeric_limits::max()) { + target.clear(); + target_width = 0; + target_height = 0; + error = "Qwen3-VL smart resize produced an unsupported image token count"; + return false; + } + image_token_count = static_cast(tokens); + return true; +} + +bool preprocess_oft_rgb(const uint8_t * source, int source_width, int source_height, + int channels, int source_stride, int training_width, + int training_height, int processor_width, int processor_height, + std::vector & target, std::string & error) { + target.clear(); + if (channels != 3) { + error = "StarVLA OFT input image must be RGB"; + return false; + } + std::vector training_image; + if (!resize_pillow_bicubic_rgb(source, source_width, source_height, source_stride, + training_width, training_height, training_image, error)) { + return false; + } + return resize_torchvision_bicubic_aa_rgb( + training_image.data(), training_width, training_height, training_width * 3, + processor_width, processor_height, target, error); +} + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/oft_image_preprocess.h b/src/models/starvla/oft_image_preprocess.h new file mode 100644 index 0000000..e4d1e02 --- /dev/null +++ b/src/models/starvla/oft_image_preprocess.h @@ -0,0 +1,34 @@ +#pragma once + +#include +#include +#include + +namespace robotcpp::starvla { + +bool resize_pillow_bicubic_rgb(const uint8_t * source, int source_width, int source_height, + int source_stride, int target_width, int target_height, + std::vector & target, std::string & error); + +bool resize_torchvision_bicubic_aa_rgb(const uint8_t * source, int source_width, + int source_height, int source_stride, int target_width, + int target_height, std::vector & target, + std::string & error); + +bool qwen3vl_smart_resize_dimensions(int source_width, int source_height, int factor, + int min_pixels, int max_pixels, int & target_width, + int & target_height, std::string & error); + +bool preprocess_qwen3vl_rgb(const uint8_t * source, int source_width, int source_height, + int channels, int source_stride, int patch_size, + int spatial_merge_size, int min_pixels, int max_pixels, + std::vector & target, int & target_width, + int & target_height, int & image_token_count, + std::string & error); + +bool preprocess_oft_rgb(const uint8_t * source, int source_width, int source_height, + int channels, int source_stride, int training_width, + int training_height, int processor_width, int processor_height, + std::vector & target, std::string & error); + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/oft_policy.cpp b/src/models/starvla/oft_policy.cpp new file mode 100644 index 0000000..48bbd62 --- /dev/null +++ b/src/models/starvla/oft_policy.cpp @@ -0,0 +1,539 @@ +#include "models/starvla/oft_policy.h" + +#include "ggml-backend.h" +#include "ggml.h" +#include "gguf.h" +#include "models/ggml_backend.h" +#include "models/gguf_loader.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace robotcpp::starvla { + +namespace { + +struct OFTBlockWeights { + ggml_tensor * norm_weight = nullptr; + ggml_tensor * norm_bias = nullptr; + ggml_tensor * linear_weight = nullptr; + ggml_tensor * linear_bias = nullptr; +}; + +struct OFTWeights { + ggml_tensor * input_norm_weight = nullptr; + ggml_tensor * input_norm_bias = nullptr; + ggml_tensor * input_proj_weight = nullptr; + ggml_tensor * input_proj_bias = nullptr; + std::vector blocks; + ggml_tensor * output_norm_weight = nullptr; + ggml_tensor * output_norm_bias = nullptr; + ggml_tensor * output_proj_weight = nullptr; + ggml_tensor * output_proj_bias = nullptr; +}; + +int require_key(gguf_context * gguf, const char * key, gguf_type type) { + const int index = gguf_find_key(gguf, key); + if (index < 0) { + throw std::runtime_error(std::string("missing required StarVLA GGUF metadata: ") + key); + } + if (gguf_get_kv_type(gguf, index) != type) { + throw std::runtime_error(std::string("invalid StarVLA GGUF metadata type: ") + key); + } + return index; +} + +std::string require_string(gguf_context * gguf, const char * key) { + return gguf_get_val_str(gguf, require_key(gguf, key, GGUF_TYPE_STRING)); +} + +int require_i32(gguf_context * gguf, const char * key) { + return gguf_get_val_i32(gguf, require_key(gguf, key, GGUF_TYPE_INT32)); +} + +float require_f32(gguf_context * gguf, const char * key) { + return gguf_get_val_f32(gguf, require_key(gguf, key, GGUF_TYPE_FLOAT32)); +} + +bool require_bool(gguf_context * gguf, const char * key) { + return gguf_get_val_bool(gguf, require_key(gguf, key, GGUF_TYPE_BOOL)); +} + +int require_array(gguf_context * gguf, const char * key, gguf_type element_type) { + const int index = require_key(gguf, key, GGUF_TYPE_ARRAY); + if (gguf_get_arr_type(gguf, index) != element_type) { + throw std::runtime_error(std::string("invalid StarVLA GGUF array element type: ") + key); + } + return index; +} + +std::vector require_string_array(gguf_context * gguf, const char * key) { + const int index = require_array(gguf, key, GGUF_TYPE_STRING); + const size_t count = gguf_get_arr_n(gguf, index); + std::vector values; + values.reserve(count); + for (size_t i = 0; i < count; ++i) { + values.emplace_back(gguf_get_arr_str(gguf, index, i)); + } + return values; +} + +std::vector require_i32_array(gguf_context * gguf, const char * key) { + const int index = require_array(gguf, key, GGUF_TYPE_INT32); + const size_t count = gguf_get_arr_n(gguf, index); + if (count == 0) { + return {}; + } + const auto * data = static_cast(gguf_get_arr_data(gguf, index)); + if (data == nullptr) { + throw std::runtime_error(std::string("missing StarVLA GGUF array data: ") + key); + } + return std::vector(data, data + count); +} + +std::vector require_f32_array(gguf_context * gguf, const char * key) { + const int index = require_array(gguf, key, GGUF_TYPE_FLOAT32); + const size_t count = gguf_get_arr_n(gguf, index); + if (count == 0) { + return {}; + } + const auto * data = static_cast(gguf_get_arr_data(gguf, index)); + if (data == nullptr) { + throw std::runtime_error(std::string("missing StarVLA GGUF array data: ") + key); + } + return std::vector(data, data + count); +} + +std::vector require_bool_array(gguf_context * gguf, const char * key) { + const int index = require_array(gguf, key, GGUF_TYPE_BOOL); + const size_t count = gguf_get_arr_n(gguf, index); + const auto * data = static_cast(gguf_get_arr_data(gguf, index)); + if (data == nullptr && count != 0) { + throw std::runtime_error(std::string("missing StarVLA GGUF array data: ") + key); + } + std::vector values(count); + for (size_t i = 0; i < count; ++i) { + values[i] = data[i] != 0 ? 1 : 0; + } + return values; +} + +std::string profile_key(int profile_index, const char * suffix) { + return "starvla.normalization.profile." + std::to_string(profile_index) + "." + suffix; +} + +bool has_shape(const ggml_tensor * tensor, std::initializer_list expected) { + if (tensor == nullptr || static_cast(ggml_n_dims(tensor)) != expected.size()) { + return false; + } + size_t dimension = 0; + for (int64_t value : expected) { + if (tensor->ne[dimension++] != value) { + return false; + } + } + return true; +} + +const char * mode_name(backend_mode mode) { + switch (mode) { + case backend_mode::cpu: + return "cpu"; + case backend_mode::cuda: + return "cuda"; + case backend_mode::metal: + return "metal"; + } + return "unknown"; +} + +class OFTGGUFLoader final : public gguf_loader { + public: + OFTGGUFLoader(OFTPolicyConfig & config, OFTWeights & weights) : config_(config), weights_(weights) {} + + protected: + bool parse_metadata(gguf_context * gguf) override { + const std::string architecture = require_string(gguf, "general.architecture"); + if (architecture != "starvla-policy") { + throw std::runtime_error("StarVLA policy GGUF has incompatible general.architecture: " + architecture); + } + if (require_i32(gguf, "starvla.schema_version") != 1 || + require_string(gguf, "starvla.framework") != "oft") { + throw std::runtime_error("StarVLA policy GGUF is not a supported Qwen OFT schema"); + } + config_.backbone_arch = require_string(gguf, "starvla.backbone.arch"); + if (config_.backbone_arch != "qwen3_vl" && + config_.backbone_arch != "qwen2_5_vl") { + throw std::runtime_error( + "StarVLA OFT policy has an unsupported Qwen backbone"); + } + + config_.bundle_uuid = require_string(gguf, "starvla.bundle.uuid"); + if (config_.bundle_uuid.empty()) { + throw std::runtime_error("StarVLA policy bundle UUID is missing"); + } + config_.text_filename = require_string(gguf, "starvla.component.text.filename"); + config_.mmproj_filename = require_string(gguf, "starvla.component.mmproj.filename"); + if (config_.text_filename.empty() || config_.mmproj_filename.empty()) { + throw std::runtime_error("StarVLA policy component filenames must be non-empty"); + } + config_.input_dim = require_i32(gguf, "starvla.qwen.hidden_size"); + config_.input_embedding_dim = require_i32(gguf, "starvla.qwen.input_embedding_size"); + config_.vocab_size = require_i32(gguf, "starvla.qwen.vocab_size"); + config_.hidden_dim = require_i32(gguf, "starvla.oft.hidden_size"); + config_.block_count = require_i32(gguf, "starvla.oft.block_count"); + config_.action_dim = require_i32(gguf, "starvla.action.dimension"); + config_.horizon = require_i32(gguf, "starvla.action.horizon"); + config_.layer_norm_epsilon = require_f32(gguf, "starvla.oft.layer_norm_epsilon"); + if (config_.input_dim <= 0 || config_.input_embedding_dim <= 0 || + config_.vocab_size <= 0 || config_.hidden_dim <= 0 || + config_.block_count <= 0 || + config_.action_dim <= 0 || config_.horizon <= 0 || !std::isfinite(config_.layer_norm_epsilon) || + config_.layer_norm_epsilon <= 0.0f) { + throw std::runtime_error("StarVLA OFT policy metadata has incompatible dimensions"); + } + + config_.prompt.horizon = config_.horizon; + config_.prompt.action_token = require_string(gguf, "starvla.prompt.action_token"); + config_.prompt.action_suffix = require_string(gguf, "starvla.prompt.action_suffix"); + config_.prompt.cot_enabled = require_bool(gguf, "starvla.prompt.cot_enabled"); + config_.prompt.cot_template = require_string(gguf, "starvla.prompt.cot_template"); + config_.prompt.state_bins = require_i32(gguf, "starvla.prompt.state_bins"); + config_.prompt.state_bin_min = require_f32(gguf, "starvla.prompt.state_bin_min"); + config_.prompt.state_bin_max = require_f32(gguf, "starvla.prompt.state_bin_max"); + config_.prompt.state_clip = require_bool(gguf, "starvla.prompt.state_clip"); + config_.action_token_id = require_i32(gguf, "starvla.prompt.action_token_id"); + std::string prompt_error; + if (!validate_oft_prompt_config(config_.prompt, prompt_error) || config_.action_token_id < 0) { + throw std::runtime_error(prompt_error.empty() ? + "StarVLA OFT token/template metadata is incompatible" : + prompt_error); + } + config_.image_count = require_i32(gguf, "starvla.image.count"); + config_.image_names = require_string_array(gguf, "starvla.image.names"); + config_.image_processor_min_pixels = + require_i32(gguf, "starvla.image.processor_min_pixels"); + config_.image_processor_max_pixels = + require_i32(gguf, "starvla.image.processor_max_pixels"); + config_.image_patch_size = require_i32(gguf, "starvla.image.patch_size"); + config_.image_spatial_merge_size = + require_i32(gguf, "starvla.image.spatial_merge_size"); + config_.image_min_token_count = + require_i32(gguf, "starvla.image.min_token_count"); + config_.image_max_token_count = + require_i32(gguf, "starvla.image.max_token_count"); + if (config_.image_count <= 0 || + config_.image_names.size() != static_cast(config_.image_count) || + config_.image_processor_min_pixels <= 0 || + config_.image_processor_max_pixels < config_.image_processor_min_pixels || + config_.image_patch_size <= 0 || + config_.image_spatial_merge_size <= 0 || + config_.image_min_token_count <= 0 || + config_.image_max_token_count < config_.image_min_token_count) { + throw std::runtime_error("StarVLA OFT image metadata is incompatible"); + } + + NormalizationConfig & normalization = config_.normalization; + normalization.clip_actions = require_bool(gguf, "starvla.normalization.clip_actions"); + normalization.binary_threshold = require_f32(gguf, "starvla.normalization.binary_threshold"); + normalization.binary_comparison = require_string(gguf, "starvla.normalization.binary_comparison"); + normalization.continuous_dimensions = + require_i32_array(gguf, "starvla.action.continuous_dimensions"); + normalization.binary_dimensions = require_i32_array(gguf, "starvla.action.binary_dimensions"); + + const int profile_count = require_i32(gguf, "starvla.normalization.profile_count"); + const std::vector keys = + require_string_array(gguf, "starvla.normalization.profile_keys"); + if (profile_count <= 0 || keys.size() != static_cast(profile_count)) { + throw std::runtime_error("StarVLA normalization profile count is inconsistent"); + } + normalization.profiles.clear(); + normalization.profiles.reserve(static_cast(profile_count)); + for (int i = 0; i < profile_count; ++i) { + NormalizationProfile profile; + const std::string key_field = profile_key(i, "key"); + const std::string q01_field = profile_key(i, "action_q01"); + const std::string q99_field = profile_key(i, "action_q99"); + const std::string mask_field = profile_key(i, "action_mask"); + profile.key = require_string(gguf, key_field.c_str()); + profile.action_q01 = require_f32_array(gguf, q01_field.c_str()); + profile.action_q99 = require_f32_array(gguf, q99_field.c_str()); + profile.action_mask = require_bool_array(gguf, mask_field.c_str()); + if (profile.key != keys[static_cast(i)]) { + throw std::runtime_error("StarVLA normalization profile order is inconsistent"); + } + normalization.profiles.push_back(std::move(profile)); + } + std::string normalization_error; + if (!validate_normalization_config(normalization, config_.action_dim, normalization_error)) { + throw std::runtime_error(normalization_error); + } + return true; + } + + bool bind_tensors(ggml_context * ctx_data) override { + weights_.input_norm_weight = require_tensor(ctx_data, "starvla.policy.oft.input_norm.weight"); + weights_.input_norm_bias = require_tensor(ctx_data, "starvla.policy.oft.input_norm.bias"); + weights_.input_proj_weight = require_tensor(ctx_data, "starvla.policy.oft.input_proj.weight"); + weights_.input_proj_bias = require_tensor(ctx_data, "starvla.policy.oft.input_proj.bias"); + weights_.blocks.clear(); + weights_.blocks.reserve(static_cast(config_.block_count)); + for (int block = 0; block < config_.block_count; ++block) { + const std::string prefix = "starvla.policy.oft.block." + std::to_string(block) + "."; + OFTBlockWeights current; + current.norm_weight = require_tensor(ctx_data, prefix + "norm.weight"); + current.norm_bias = require_tensor(ctx_data, prefix + "norm.bias"); + current.linear_weight = require_tensor(ctx_data, prefix + "linear.weight"); + current.linear_bias = require_tensor(ctx_data, prefix + "linear.bias"); + weights_.blocks.push_back(current); + } + weights_.output_norm_weight = require_tensor(ctx_data, "starvla.policy.oft.output_norm.weight"); + weights_.output_norm_bias = require_tensor(ctx_data, "starvla.policy.oft.output_norm.bias"); + weights_.output_proj_weight = require_tensor(ctx_data, "starvla.policy.oft.output_proj.weight"); + weights_.output_proj_bias = require_tensor(ctx_data, "starvla.policy.oft.output_proj.bias"); + + if (!has_shape(weights_.input_norm_weight, {config_.input_dim}) || + !has_shape(weights_.input_norm_bias, {config_.input_dim}) || + !has_shape(weights_.input_proj_weight, {config_.input_dim, config_.hidden_dim}) || + !has_shape(weights_.input_proj_bias, {config_.hidden_dim}) || + !has_shape(weights_.output_norm_weight, {config_.hidden_dim}) || + !has_shape(weights_.output_norm_bias, {config_.hidden_dim}) || + !has_shape(weights_.output_proj_weight, {config_.hidden_dim, config_.action_dim}) || + !has_shape(weights_.output_proj_bias, {config_.action_dim})) { + throw std::runtime_error("StarVLA OFT projection tensor has an incompatible ggml shape"); + } + for (const OFTBlockWeights & block : weights_.blocks) { + if (!has_shape(block.norm_weight, {config_.hidden_dim}) || + !has_shape(block.norm_bias, {config_.hidden_dim}) || + !has_shape(block.linear_weight, {config_.hidden_dim, config_.hidden_dim}) || + !has_shape(block.linear_bias, {config_.hidden_dim})) { + throw std::runtime_error("StarVLA OFT residual block tensor has an incompatible ggml shape"); + } + } + return true; + } + + private: + OFTPolicyConfig & config_; + OFTWeights & weights_; +}; + +} // namespace + +struct OFTPolicy::Impl { + OFTPolicyConfig config; + OFTWeights weights; + gguf_load_result loaded; + ggml_backend_t backend_cpu = nullptr; + std::vector backends; + ggml_backend_sched_t scheduler = nullptr; + backend_buft_policy buft_policy; + backend_mode mode = backend_mode::cpu; + int n_threads = 0; + int verbosity = 0; + ggml_context * graph_context = nullptr; + ggml_cgraph * graph = nullptr; + ggml_tensor * input = nullptr; + ggml_tensor * output = nullptr; + + ~Impl() { + if (scheduler != nullptr) { + ggml_backend_sched_synchronize(scheduler); + ggml_backend_sched_free(scheduler); + scheduler = nullptr; + } + if (graph_context != nullptr) { + ggml_free(graph_context); + graph_context = nullptr; + } + if (loaded.model_buffer != nullptr) { + ggml_backend_buffer_free(loaded.model_buffer); + loaded.model_buffer = nullptr; + } + if (loaded.ctx_data != nullptr) { + ggml_free(loaded.ctx_data); + loaded.ctx_data = nullptr; + } + if (loaded.gguf != nullptr) { + gguf_free(loaded.gguf); + loaded.gguf = nullptr; + } + for (ggml_backend_t backend : backends) { + if (backend != nullptr) { + ggml_backend_free(backend); + } + } + backends.clear(); + backend_cpu = nullptr; + } + + void build_graph() { + const size_t graph_size = GGML_DEFAULT_GRAPH_SIZE; + ggml_init_params params{}; + params.mem_size = graph_size * ggml_tensor_overhead() + ggml_graph_overhead_custom(graph_size, false); + params.mem_buffer = nullptr; + params.no_alloc = true; + graph_context = ggml_init(params); + if (graph_context == nullptr) { + throw std::runtime_error("failed to initialize StarVLA OFT graph context"); + } + + auto f32_vector = [&](ggml_tensor * tensor) { + return tensor->type == GGML_TYPE_F32 ? tensor : ggml_cast(graph_context, tensor, GGML_TYPE_F32); + }; + auto layer_norm = [&](ggml_tensor * value, ggml_tensor * weight, ggml_tensor * bias) { + ggml_tensor * normalized = ggml_norm(graph_context, value, config.layer_norm_epsilon); + normalized = ggml_mul(graph_context, normalized, f32_vector(weight)); + return ggml_add(graph_context, normalized, f32_vector(bias)); + }; + auto linear = [&](ggml_tensor * value, ggml_tensor * weight, ggml_tensor * bias) { + ggml_tensor * projected = ggml_mul_mat(graph_context, weight, value); + ggml_mul_mat_set_prec(projected, GGML_PREC_F32); + return ggml_add(graph_context, projected, f32_vector(bias)); + }; + + input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, config.input_dim, config.horizon); + ggml_set_name(input, "starvla_oft_action_queries"); + ggml_set_input(input); + + ggml_tensor * current = layer_norm(input, weights.input_norm_weight, weights.input_norm_bias); + current = ggml_relu(graph_context, linear(current, weights.input_proj_weight, weights.input_proj_bias)); + for (const OFTBlockWeights & block : weights.blocks) { + ggml_tensor * residual = current; + current = layer_norm(current, block.norm_weight, block.norm_bias); + current = ggml_relu(graph_context, linear(current, block.linear_weight, block.linear_bias)); + current = ggml_add(graph_context, current, residual); + } + current = layer_norm(current, weights.output_norm_weight, weights.output_norm_bias); + output = linear(current, weights.output_proj_weight, weights.output_proj_bias); + ggml_set_name(output, "starvla_oft_normalized_actions"); + ggml_set_output(output); + + graph = ggml_new_graph_custom(graph_context, graph_size, false); + if (graph == nullptr) { + throw std::runtime_error("failed to create StarVLA OFT graph"); + } + ggml_build_forward_expand(graph, output); + ggml_backend_sched_reset(scheduler); + if (!ggml_backend_sched_alloc_graph(scheduler, graph)) { + throw std::runtime_error("failed to allocate StarVLA OFT graph"); + } + } +}; + +OFTPolicy::OFTPolicy(std::unique_ptr impl) : impl_(std::move(impl)) {} + +OFTPolicy::~OFTPolicy() = default; + +std::unique_ptr OFTPolicy::load(const std::string & path, int n_threads, int verbosity, + std::string & error) { + error.clear(); + if (path.empty()) { + error = "StarVLA OFT policy path is required"; + return nullptr; + } + + std::unique_ptr impl(new Impl()); + impl->n_threads = n_threads; + impl->verbosity = verbosity; + try { + backend_scheduler_config scheduler_config; + scheduler_config.max_nodes = GGML_DEFAULT_GRAPH_SIZE; + scheduler_config.parallel = false; + scheduler_config.op_offload = true; + backend_loader backend; + if (!backend.load(impl->backend_cpu, impl->backends, impl->scheduler, impl->buft_policy, true, + scheduler_config, verbosity)) { + error = "failed to initialize StarVLA OFT backend: " + backend.error(); + return nullptr; + } + impl->mode = backend.mode(); + + OFTGGUFLoader loader(impl->config, impl->weights); + if (!loader.load(path.c_str(), impl->buft_policy.model_buft, impl->loaded, verbosity)) { + error = loader.error(); + return nullptr; + } + if (impl->loaded.ctx_data == nullptr || impl->loaded.model_buffer == nullptr) { + error = "StarVLA OFT policy GGUF has no tensors"; + return nullptr; + } + ggml_backend_buffer_set_usage(impl->loaded.model_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + impl->build_graph(); + if (verbosity >= 1) { + std::fprintf(stderr, + "%s: backend=%s input=%d hidden=%d blocks=%d horizon=%d action_dim=%d profiles=%zu\n", + __func__, mode_name(impl->mode), impl->config.input_dim, impl->config.hidden_dim, + impl->config.block_count, impl->config.horizon, impl->config.action_dim, + impl->config.normalization.profiles.size()); + } + } catch (const std::exception & exception) { + error = exception.what(); + return nullptr; + } + return std::unique_ptr(new OFTPolicy(std::move(impl))); +} + +bool OFTPolicy::evaluate(const float * action_queries, size_t element_count, + std::vector & normalized_actions, std::string & error) { + normalized_actions.clear(); + error.clear(); + if (impl_ == nullptr) { + error = "StarVLA OFT policy is not initialized"; + return false; + } + const size_t expected = static_cast(impl_->config.horizon) * + static_cast(impl_->config.input_dim); + if (action_queries == nullptr || element_count != expected) { + error = "StarVLA OFT action-query tensor has an incompatible shape"; + return false; + } + + if (impl_->scheduler == nullptr || impl_->graph == nullptr || impl_->input == nullptr || + impl_->output == nullptr) { + error = "StarVLA OFT policy graph is not initialized"; + return false; + } + + ggml_backend_tensor_set(impl_->input, action_queries, 0, element_count * sizeof(float)); + set_backend_threads(impl_->backends, impl_->n_threads); + if (ggml_backend_sched_graph_compute(impl_->scheduler, impl_->graph) != GGML_STATUS_SUCCESS) { + error = "StarVLA OFT graph compute failed"; + return false; + } + + const size_t output_count = static_cast(impl_->config.horizon) * + static_cast(impl_->config.action_dim); + normalized_actions.resize(output_count); + ggml_backend_tensor_get(impl_->output, normalized_actions.data(), 0, output_count * sizeof(float)); + return true; +} + +bool OFTPolicy::unnormalize(const std::vector & normalized_actions, const std::string & profile_key, + std::vector & actions, std::string & error) const { + if (impl_ == nullptr) { + actions.clear(); + error = "StarVLA OFT policy is not initialized"; + return false; + } + return denormalize_actions(impl_->config.normalization, profile_key, normalized_actions, impl_->config.horizon, + impl_->config.action_dim, actions, error); +} + +const OFTPolicyConfig & OFTPolicy::config() const { + if (impl_ == nullptr) { + throw std::runtime_error("StarVLA OFT policy is not initialized"); + } + return impl_->config; +} + +const char * OFTPolicy::backend_name() const { + return impl_ != nullptr ? mode_name(impl_->mode) : "unknown"; +} + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/oft_policy.h b/src/models/starvla/oft_policy.h new file mode 100644 index 0000000..44a1aed --- /dev/null +++ b/src/models/starvla/oft_policy.h @@ -0,0 +1,65 @@ +#pragma once + +#include "models/starvla/normalization.h" +#include "models/starvla/oft_prompt.h" + +#include +#include +#include +#include + +namespace robotcpp::starvla { + +struct OFTPolicyConfig { + std::string backbone_arch; + std::string bundle_uuid; + std::string text_filename; + std::string mmproj_filename; + int input_dim = 0; + int input_embedding_dim = 0; + int vocab_size = 0; + int hidden_dim = 0; + int block_count = 0; + int action_dim = 0; + int horizon = 0; + float layer_norm_epsilon = 0.0f; + OFTPromptConfig prompt; + int action_token_id = 0; + int image_count = 0; + std::vector image_names; + int image_processor_min_pixels = 0; + int image_processor_max_pixels = 0; + int image_patch_size = 0; + int image_spatial_merge_size = 0; + int image_min_token_count = 0; + int image_max_token_count = 0; + NormalizationConfig normalization; +}; + +class OFTPolicy { + public: + ~OFTPolicy(); + + OFTPolicy(const OFTPolicy &) = delete; + OFTPolicy & operator=(const OFTPolicy &) = delete; + + static std::unique_ptr load(const std::string & path, int n_threads, int verbosity, + std::string & error); + + bool evaluate(const float * action_queries, size_t element_count, std::vector & normalized_actions, + std::string & error); + bool unnormalize(const std::vector & normalized_actions, const std::string & profile_key, + std::vector & actions, std::string & error) const; + + const OFTPolicyConfig & config() const; + const char * backend_name() const; + + private: + struct Impl; + + explicit OFTPolicy(std::unique_ptr impl); + + std::unique_ptr impl_; +}; + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/oft_prompt.cpp b/src/models/starvla/oft_prompt.cpp new file mode 100644 index 0000000..28fb174 --- /dev/null +++ b/src/models/starvla/oft_prompt.cpp @@ -0,0 +1,165 @@ +#include "models/starvla/oft_prompt.h" + +#include +#include +#include +#include + +namespace robotcpp::starvla { +namespace { + +constexpr const char * kInstructionPlaceholder = "{instruction}"; +constexpr const char * kMtmdMediaMarker = "<__media__>"; + +std::string repeat(const std::string & value, int count) { + std::string result; + result.reserve(value.size() * static_cast(count)); + for (int i = 0; i < count; ++i) { + result += value; + } + return result; +} + +void replace_all(std::string & value, const std::string & needle, const std::string & replacement) { + size_t offset = 0; + while ((offset = value.find(needle, offset)) != std::string::npos) { + value.replace(offset, needle.size(), replacement); + offset += replacement.size(); + } +} + +bool discretize_state(const OFTPromptConfig & config, const std::vector & state, + std::string & output, std::string & error) { + if (state.empty()) { + output.clear(); + return true; + } + + std::ostringstream stream; + const double minimum = static_cast(config.state_bin_min); + const double maximum = static_cast(config.state_bin_max); + const double step = (maximum - minimum) / static_cast(config.state_bins); + for (size_t i = 0; i < state.size(); ++i) { + double value = static_cast(state[i]); + if (!std::isfinite(value)) { + error = "StarVLA OFT state contains a non-finite value"; + return false; + } + if (config.state_clip) { + value = std::max(minimum, std::min(maximum, value)); + } + + // Matches numpy.digitize(value, linspace(min, max, bins + 1)[:-1]) - 1. + int bin = -1; + for (int edge = 0; edge < config.state_bins; ++edge) { + const double boundary = minimum + step * static_cast(edge); + if (value >= boundary) { + bin = edge; + } else { + break; + } + } + if (i != 0) { + stream << ' '; + } + stream << bin; + } + output = stream.str(); + return true; +} + +} // namespace + +bool validate_oft_prompt_config(const OFTPromptConfig & config, std::string & error) { + error.clear(); + if (config.horizon <= 0 || config.action_token.empty()) { + error = "StarVLA OFT prompt has an invalid horizon or action token"; + return false; + } + const std::string expected_suffix = " Please predict the next " + std::to_string(config.horizon) + + " robot actions: " + + repeat(config.action_token, config.horizon) + "."; + if (config.action_suffix != expected_suffix) { + error = "StarVLA OFT action suffix does not match its horizon/token contract"; + return false; + } + if (config.cot_enabled && config.cot_template.find(kInstructionPlaceholder) == std::string::npos) { + error = "StarVLA OFT CoT template is missing {instruction}"; + return false; + } + if (config.state_bins <= 0 || + !std::isfinite(config.state_bin_min) || !std::isfinite(config.state_bin_max) || + config.state_bin_max <= config.state_bin_min) { + error = "StarVLA OFT state prompt metadata is incompatible"; + return false; + } + return true; +} + +bool build_oft_instruction(const OFTPromptConfig & config, const std::string & task, + const std::vector & state, std::string & instruction, + std::string & error) { + instruction.clear(); + if (!validate_oft_prompt_config(config, error)) { + return false; + } + if (task.find(kMtmdMediaMarker) != std::string::npos) { + error = "StarVLA OFT task contains the reserved mtmd media marker"; + return false; + } + + instruction = task; + if (!state.empty()) { + std::string state_text; + if (!discretize_state(config, state, state_text, error)) { + instruction.clear(); + return false; + } + instruction += " [STATE] " + state_text + " [ACTION]"; + } + instruction += config.action_suffix; + + if (config.cot_enabled) { + std::string wrapped = config.cot_template; + replace_all(wrapped, kInstructionPlaceholder, instruction); + instruction = std::move(wrapped); + } + return true; +} + +std::string build_qwen_media_content(size_t image_count, const std::string & instruction, + const char * media_marker) { + const std::string marker = media_marker == nullptr ? std::string() : std::string(media_marker); + std::string content; + content.reserve(marker.size() * image_count + instruction.size()); + for (size_t i = 0; i < image_count; ++i) { + content += marker; + } + content += instruction; + return content; +} + +bool find_last_token_positions(const std::vector & token_ids, int32_t token_id, + size_t count, std::vector & positions, + std::string & error) { + positions.clear(); + error.clear(); + if (count == 0) { + error = "StarVLA OFT action-token count must be positive"; + return false; + } + for (size_t i = 0; i < token_ids.size(); ++i) { + if (token_ids[i] == token_id) { + positions.push_back(i); + } + } + if (positions.size() < count) { + error = "StarVLA OFT prompt contains fewer action tokens than its horizon"; + positions.clear(); + return false; + } + positions.erase(positions.begin(), positions.end() - static_cast(count)); + return true; +} + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/oft_prompt.h b/src/models/starvla/oft_prompt.h new file mode 100644 index 0000000..4f70435 --- /dev/null +++ b/src/models/starvla/oft_prompt.h @@ -0,0 +1,35 @@ +#pragma once + +#include +#include +#include +#include + +namespace robotcpp::starvla { + +struct OFTPromptConfig { + int horizon = 0; + std::string action_token; + std::string action_suffix; + bool cot_enabled = false; + std::string cot_template; + int state_bins = 0; + float state_bin_min = 0.0f; + float state_bin_max = 0.0f; + bool state_clip = false; +}; + +bool validate_oft_prompt_config(const OFTPromptConfig & config, std::string & error); + +bool build_oft_instruction(const OFTPromptConfig & config, const std::string & task, + const std::vector & state, std::string & instruction, + std::string & error); + +std::string build_qwen_media_content(size_t image_count, const std::string & instruction, + const char * media_marker); + +bool find_last_token_positions(const std::vector & token_ids, int32_t token_id, + size_t count, std::vector & positions, + std::string & error); + +} // namespace robotcpp::starvla diff --git a/tests/starvla/oft_image_preprocess_test.cpp b/tests/starvla/oft_image_preprocess_test.cpp new file mode 100644 index 0000000..89f7d95 --- /dev/null +++ b/tests/starvla/oft_image_preprocess_test.cpp @@ -0,0 +1,162 @@ +#include "models/starvla/oft_image_preprocess.h" + +#include +#include +#include +#include +#include + +namespace { + +void require(bool condition, const char * message) { + if (!condition) { + std::cerr << "FAIL: " << message << '\n'; + std::exit(1); + } +} + +uint64_t fnv1a(const std::vector & values) { + uint64_t hash = UINT64_C(14695981039346656037); + for (uint8_t value : values) { + hash ^= value; + hash *= UINT64_C(1099511628211); + } + return hash; +} + +std::vector test_image(int width, int height, int stride) { + std::vector image(static_cast(stride) * height, 0xee); + for (int y = 0; y < height; ++y) { + for (int x = 0; x < width; ++x) { + for (int c = 0; c < 3; ++c) { + image[static_cast(y) * stride + x * 3 + c] = + static_cast((x * 37 + y * 61 + c * 83 + x * y * 7) & 0xff); + } + } + } + return image; +} + +} // namespace + +int main() { + using namespace robotcpp::starvla; + + const int width = 7; + const int height = 5; + const int stride = width * 3 + 5; + const std::vector source = test_image(width, height, stride); + std::vector pillow; + std::vector torch; + std::string error; + require(resize_pillow_bicubic_rgb(source.data(), width, height, stride, 4, 6, pillow, error), + "Pillow bicubic resize must succeed"); + require(resize_torchvision_bicubic_aa_rgb(source.data(), width, height, stride, 8, 9, torch, error), + "torchvision bicubic resize must succeed"); + require(pillow.size() == 4 * 6 * 3, "Pillow resize shape must be exact"); + require(torch.size() == 8 * 9 * 3, "torchvision resize shape must be exact"); + + // Byte-for-byte references: Pillow 7.0.0 and torchvision's uint8 bicubic-AA path. + require(fnv1a(pillow) == UINT64_C(13888783175115780895), + "Pillow bicubic resize must match the reference bytes"); + require(fnv1a(torch) == UINT64_C(14284510537394890340), + "torchvision bicubic-AA resize must match the reference bytes"); + + std::vector processed; + require(preprocess_oft_rgb(source.data(), width, height, 3, stride, 4, 6, 8, 9, + processed, error), + "two-stage OFT preprocessing must succeed"); + require(processed.size() == 8 * 9 * 3, "two-stage OFT image shape must be exact"); + require(fnv1a(processed) == UINT64_C(11011422699469149164), + "two-stage OFT preprocessing must match the reference bytes"); + + int smart_width = 0; + int smart_height = 0; + require(qwen3vl_smart_resize_dimensions(640, 488, 32, 65536, 16777216, + smart_width, smart_height, error), + "Qwen3-VL smart resize dimensions must succeed"); + require(smart_width == 640 && smart_height == 480, + "640x488 must map to the official 640x480 grid"); + require(qwen3vl_smart_resize_dimensions(640, 400, 32, 65536, 16777216, + smart_width, smart_height, error), + "Qwen3-VL ties-to-even smart resize must succeed"); + require(smart_width == 640 && smart_height == 384, + "Python ties-to-even rounding must map 400 pixels to 384"); + + std::vector smart; + int smart_tokens = 0; + require(preprocess_qwen3vl_rgb(source.data(), width, height, 3, stride, 16, 2, + 65536, 16777216, smart, smart_width, smart_height, + smart_tokens, error), + "Qwen3-VL dynamic preprocessing must succeed"); + require(smart_width == 320 && smart_height == 224 && smart_tokens == 70, + "Qwen3-VL dynamic preprocessing must report its exact grid and token count"); + require(smart.size() == static_cast(smart_width) * smart_height * 3, + "Qwen3-VL dynamic image shape must be exact"); + + const int official_width = 640; + const int official_height = 488; + const int official_stride = official_width * 3 + 5; + const std::vector official_shape_source = + test_image(official_width, official_height, official_stride); + require(preprocess_qwen3vl_rgb( + official_shape_source.data(), official_width, official_height, 3, + official_stride, 16, 2, 65536, 16777216, smart, smart_width, + smart_height, smart_tokens, error), + "official-shape Qwen3-VL preprocessing must succeed"); + require(smart_width == 640 && smart_height == 480 && smart_tokens == 300, + "official-shape Qwen3-VL preprocessing must produce a 30x40 merged grid"); + // Reference generated by torchvision 0.21's uint8 bicubic antialias path. + require(fnv1a(smart) == UINT64_C(11336294493867056015), + "official-shape Qwen3-VL resize must match torchvision bytes"); + + require(preprocess_qwen3vl_rgb( + official_shape_source.data(), official_width, official_height, 3, + official_stride, 14, 2, 3136, 12845056, smart, smart_width, + smart_height, smart_tokens, error), + "official-shape Qwen2.5-VL preprocessing must succeed"); + require(smart_width == 644 && smart_height == 476 && smart_tokens == 391, + "Qwen2.5-VL preprocessing must produce its exact 17x23 merged grid"); + // Reference generated by Transformers 4.57 / torchvision 0.21 fast Qwen2-VL + // image preprocessing on the same uint8 RGB tensor. + require(fnv1a(smart) == UINT64_C(12050652900109057577), + "Qwen2.5-VL resize must match torchvision bytes"); + + const int bridge_width = 256; + const int bridge_height = 256; + const int bridge_stride = bridge_width * 3; + const std::vector bridge_source = + test_image(bridge_width, bridge_height, bridge_stride); + require(preprocess_qwen3vl_rgb( + bridge_source.data(), bridge_width, bridge_height, 3, + bridge_stride, 14, 2, 3136, 12845056, smart, smart_width, + smart_height, smart_tokens, error), + "Bridge-size Qwen2.5-VL preprocessing must succeed"); + require(smart_width == 252 && smart_height == 252 && smart_tokens == 81, + "256x256 Bridge input must map to the official 18x18 patch grid"); + require(smart.size() == static_cast(252 * 252 * 3), + "Bridge-size Qwen2.5-VL preprocessing must return 252x252 RGB"); + + const int adapter_width = 224; + const int adapter_height = 224; + const int adapter_stride = adapter_width * 3; + const std::vector adapter_source = + test_image(adapter_width, adapter_height, adapter_stride); + require(preprocess_qwen3vl_rgb( + adapter_source.data(), adapter_width, adapter_height, 3, + adapter_stride, 14, 2, 3136, 12845056, smart, smart_width, + smart_height, smart_tokens, error), + "SimplerEnv adapter-size Qwen2.5-VL preprocessing must succeed"); + require(smart_width == 224 && smart_height == 224 && smart_tokens == 64, + "224x224 SimplerEnv input must preserve its official 16x16 patch grid"); + require(smart.size() == static_cast(224 * 224 * 3), + "SimplerEnv adapter-size preprocessing must return 224x224 RGB"); + + require(!qwen3vl_smart_resize_dimensions(640, 3, 32, 65536, 16777216, + smart_width, smart_height, error), + "Qwen3-VL smart resize must reject aspect ratios over 200"); + + std::cout << "pillow_fnv=" << fnv1a(pillow) << " torch_fnv=" << fnv1a(torch) + << " two_stage_fnv=" << fnv1a(processed) << '\n'; + return 0; +} diff --git a/tests/starvla/oft_prompt_test.cpp b/tests/starvla/oft_prompt_test.cpp new file mode 100644 index 0000000..5c6b838 --- /dev/null +++ b/tests/starvla/oft_prompt_test.cpp @@ -0,0 +1,81 @@ +#include "models/starvla/oft_prompt.h" + +#include +#include +#include +#include + +namespace { + +void require(bool condition, const char * message) { + if (!condition) { + std::cerr << "FAIL: " << message << '\n'; + std::exit(1); + } +} + +robotcpp::starvla::OFTPromptConfig official_config() { + robotcpp::starvla::OFTPromptConfig config; + config.horizon = 16; + config.action_token = "\xF0\x9F\x94\x8D"; + config.action_suffix = " Please predict the next 16 robot actions: "; + for (int i = 0; i < config.horizon; ++i) { + config.action_suffix += config.action_token; + } + config.action_suffix += "."; + config.cot_enabled = true; + config.cot_template = "Your task is {instruction}. Locate {instruction}."; + config.state_bins = 256; + config.state_bin_min = -1.0f; + config.state_bin_max = 1.0f; + config.state_clip = false; + return config; +} + +} // namespace + +int main() { + using namespace robotcpp::starvla; + + OFTPromptConfig config = official_config(); + std::string error; + require(validate_oft_prompt_config(config, error), "official prompt config must validate"); + + config.cot_enabled = false; + std::string instruction; + require(build_oft_instruction(config, "grab", {}, instruction, error), + "prompt without state must build"); + require(instruction == "grab" + config.action_suffix, "action suffix placement must match StarVLA"); + require(!build_oft_instruction(config, "grab <__media__> now", {}, instruction, error), + "reserved mtmd media markers in tasks must be rejected"); + require(error.find("reserved mtmd media marker") != std::string::npos, + "media marker rejection must explain the contract violation"); + + require(build_oft_instruction(config, "grab", {-1.1f, -1.0f, 0.0f, 1.0f, 1.1f}, + instruction, error), + "state prompt must build"); + require(instruction == "grab [STATE] -1 0 128 255 255 [ACTION]" + config.action_suffix, + "state discretization must match numpy.digitize"); + + config.cot_enabled = true; + require(build_oft_instruction(config, "grab", {}, instruction, error), "CoT prompt must build"); + const std::string unwrapped = "grab" + config.action_suffix; + require(instruction == "Your task is " + unwrapped + ". Locate " + unwrapped + ".", + "Python str.replace semantics must replace every instruction placeholder"); + + const std::string content = build_qwen_media_content(2, instruction, "<__media__>"); + require(content == "<__media__><__media__>" + instruction, + "image markers must precede text without separators"); + + std::vector ids = {1, 9, 2, 9, 9, 3, 9}; + std::vector positions; + require(find_last_token_positions(ids, 9, 3, positions, error), "action tokens must be found"); + require(positions == std::vector({3, 4, 6}), + "last action positions must remain in temporal order"); + require(!find_last_token_positions(ids, 9, 5, positions, error), + "insufficient action tokens must fail"); + require(!error.empty(), "insufficient action tokens must return an error"); + + std::cout << "starvla OFT prompt tests passed\n"; + return 0; +} diff --git a/tools/hf2gguf/starvla/generate_starvla_oft_golden.py b/tools/hf2gguf/starvla/generate_starvla_oft_golden.py new file mode 100644 index 0000000..bf9095f --- /dev/null +++ b/tools/hf2gguf/starvla/generate_starvla_oft_golden.py @@ -0,0 +1,928 @@ +#!/usr/bin/env python3 +"""Generate an auditable Python oracle from the pinned official StarVLA OFT checkpoint. + +This intentionally executes StarVLA's own Qwenvl_OFT preprocessing, Qwen3-VL +forward, action-token gather, OFT head, and PolicyNormProcessor. The emitted +JSON/NPZ pair is the value-level reference used to qualify GGUF inference. +""" + +from __future__ import annotations + +import argparse +import contextlib +import datetime as dt +import gc +import hashlib +import importlib.metadata +import json +import os +import platform +import shutil +import subprocess +import sys +import tempfile +from pathlib import Path +from typing import Any, Iterable, Mapping, Sequence + +import numpy as np + + +TOOLS_DIR = Path(__file__).resolve().parent +if str(TOOLS_DIR) not in sys.path: + sys.path.insert(0, str(TOOLS_DIR)) + +from starvla_checkpoint import ( # noqa: E402 + DEFAULT_CATALOG, + StarVLAError, + get_variant, + load_catalog, + official_bundle_uuid, + sha256_file, + verify_catalog_files, + verify_checkpoint_file, +) + + +GOLDEN_SCHEMA_VERSION = 1 +SUPPORTED_VARIANT = "oft" +ACTION_TOKEN = chr(0x1F50D) +ACTION_TOKEN_ID = 146663 +EXPECTED_TRANSFORMERS_VERSION = "4.57.0" +EXPECTED_TORCH_VERSION = "2.6.0" +EXPECTED_TORCHVISION_VERSION = "0.21.0" +EXPECTED_NUMPY_VERSION = "1.26.4" + + +def _canonical_json(value: Any) -> bytes: + return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode("utf-8") + + +def _sha256_bytes(value: bytes) -> str: + return hashlib.sha256(value).hexdigest() + + +def _array_sha256(value: np.ndarray) -> str: + array = np.ascontiguousarray(value) + header = _canonical_json({"dtype": array.dtype.str, "shape": list(array.shape)}) + return _sha256_bytes(header + b"\x00" + array.tobytes(order="C")) + + +def _array_record(value: np.ndarray, *, source_dtype: str | None = None) -> dict[str, Any]: + array = np.ascontiguousarray(value) + record: dict[str, Any] = { + "dtype": array.dtype.str, + "shape": list(array.shape), + "sha256": _array_sha256(array), + } + if source_dtype is not None: + record["source_dtype"] = source_dtype + return record + + +def _distribution_version(name: str) -> str: + try: + return importlib.metadata.version(name) + except importlib.metadata.PackageNotFoundError: + return "missing" + + +def _base_version(version: str) -> str: + return version.split("+", 1)[0] + + +def validate_runtime_versions( + *, + torch_version: str, + torchvision_version: str, + transformers_version: str, + numpy_version: str, +) -> None: + expected = { + "torch": EXPECTED_TORCH_VERSION, + "torchvision": EXPECTED_TORCHVISION_VERSION, + "transformers": EXPECTED_TRANSFORMERS_VERSION, + "numpy": EXPECTED_NUMPY_VERSION, + } + actual = { + "torch": _base_version(torch_version), + "torchvision": _base_version(torchvision_version), + "transformers": _base_version(transformers_version), + "numpy": _base_version(numpy_version), + } + mismatches = [ + f"{name}: expected {expected[name]}, got {actual[name]}" + for name in expected + if actual[name] != expected[name] + ] + if mismatches: + raise StarVLAError("official oracle runtime version mismatch: " + "; ".join(mismatches)) + + +def select_action_positions( + input_ids: np.ndarray, + *, + action_token_id: int, + chunk_len: int, +) -> tuple[list[list[int]], list[list[int]]]: + ids = np.asarray(input_ids) + if ids.ndim != 2: + raise StarVLAError(f"input_ids must be rank 2, got {list(ids.shape)}") + if chunk_len <= 0: + raise StarVLAError(f"action chunk length must be positive, got {chunk_len}") + + all_positions: list[list[int]] = [] + selected_positions: list[list[int]] = [] + for batch_index, row in enumerate(ids): + positions = np.flatnonzero(row == action_token_id).astype(np.int64).tolist() + if len(positions) < chunk_len: + raise StarVLAError( + f"sample {batch_index} has {len(positions)} action tokens; expected at least {chunk_len}" + ) + all_positions.append(positions) + selected_positions.append(positions[-chunk_len:]) + return all_positions, selected_positions + + +def expected_framework_instruction(config: Mapping[str, Any], task: str, chunk_len: int) -> str: + if not isinstance(task, str) or not task.strip(): + raise StarVLAError("task must be a non-empty string") + try: + vla_data = config["datasets"]["vla_data"] + except (KeyError, TypeError) as exc: + raise StarVLAError("checkpoint config has no datasets.vla_data object") from exc + if not isinstance(vla_data, Mapping): + raise StarVLAError("checkpoint config datasets.vla_data must be an object") + action_tokens = ACTION_TOKEN * chunk_len + return task + f" Please predict the next {chunk_len} robot actions: {action_tokens}." + + +def expected_model_instruction(config: Mapping[str, Any], framework_instruction: str) -> str: + """Mirror QWen3.build_qwenvl_inputs after QwenOFT adds its action suffix.""" + + try: + vla_data = config["datasets"]["vla_data"] + except (KeyError, TypeError) as exc: + raise StarVLAError("checkpoint config has no datasets.vla_data object") from exc + if not isinstance(vla_data, Mapping): + raise StarVLAError("checkpoint config datasets.vla_data must be an object") + cot_prompt = vla_data.get("CoT_prompt") + return ( + cot_prompt.replace("{instruction}", framework_instruction) + if isinstance(cot_prompt, str) + else framework_instruction + ) + + +def _run_git(source_dir: Path, *arguments: str) -> str: + try: + result = subprocess.run( + ["git", "-C", str(source_dir), *arguments], + check=True, + capture_output=True, + text=True, + ) + except (OSError, subprocess.CalledProcessError) as exc: + raise StarVLAError(f"failed to inspect pinned StarVLA checkout {source_dir}: {exc}") from exc + return result.stdout.strip() + + +def verify_pinned_source_checkout(source_dir: Path, expected_revision: str) -> None: + source_dir = source_dir.resolve() + if not (source_dir / ".git").exists(): + raise StarVLAError(f"StarVLA source is not a Git checkout: {source_dir}") + actual_revision = _run_git(source_dir, "rev-parse", "HEAD") + if actual_revision != expected_revision: + raise StarVLAError( + f"StarVLA source revision mismatch: expected {expected_revision}, got {actual_revision}" + ) + changes = _run_git(source_dir, "status", "--porcelain=v1", "--untracked-files=all") + if changes: + raise StarVLAError(f"pinned StarVLA checkout has tracked or untracked changes:\n{changes}") + + +def _ensure_regular_file(path: Path, *, label: str) -> None: + if not path.is_file() or path.is_symlink(): + raise StarVLAError(f"{label} must be a regular, non-symlink file: {path}") + + +def validate_official_inputs( + *, + checkpoint_root: Path, + source_dir: Path, + catalog_path: Path = DEFAULT_CATALOG, +) -> dict[str, Any]: + catalog = load_catalog(catalog_path) + variant = get_variant(catalog, SUPPORTED_VARIANT) + qwen = catalog["shared_assets"]["qwen3_vl_4b_instruct"] + checkpoint_root = checkpoint_root.resolve() + policy_dir = checkpoint_root / "sources" / variant["directory"] / variant["revision"] + qwen_dir = checkpoint_root / "sources" / qwen["directory"] / qwen["revision"] + checkpoint = policy_dir / variant["checkpoint"]["path"] + + expected_source = (checkpoint_root / "source" / "starvla").resolve() + if source_dir.resolve() != expected_source: + raise StarVLAError( + f"StarVLA source must be the canonical checkout {expected_source}, got {source_dir.resolve()}" + ) + verify_pinned_source_checkout(source_dir, catalog["source_revisions"]["starvla"]) + verify_catalog_files(policy_dir, variant) + verify_catalog_files(qwen_dir, qwen) + _ensure_regular_file(checkpoint, label="official OFT checkpoint") + incomplete_sidecar = Path(f"{checkpoint}.aria2") + if incomplete_sidecar.exists(): + raise StarVLAError( + f"official OFT checkpoint download is incomplete ({incomplete_sidecar} exists); resume the download first" + ) + verify_checkpoint_file(checkpoint, variant) + return { + "catalog": catalog, + "variant": variant, + "qwen": qwen, + "policy_dir": policy_dir, + "qwen_dir": qwen_dir, + "checkpoint": checkpoint, + "source_dir": source_dir.resolve(), + "catalog_path": catalog_path.resolve(), + } + + +def _require_isolated_python() -> None: + if not sys.flags.isolated: + raise StarVLAError( + "the oracle must run in Python isolated mode; invoke it as `python -I " + "tools/hf2gguf/starvla/generate_starvla_oft_golden.py ...`" + ) + + +def _configure_determinism(torch: Any, *, seed: int, device: str) -> None: + os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" + if not device.startswith("cuda"): + raise StarVLAError("the official OFT golden oracle currently requires a CUDA device") + if not torch.cuda.is_available(): + raise StarVLAError("CUDA is not available to PyTorch") + try: + device_index = torch.device(device).index + except (RuntimeError, ValueError) as exc: + raise StarVLAError(f"invalid CUDA device {device!r}: {exc}") from exc + torch.cuda.set_device(0 if device_index is None else device_index) + if not torch.cuda.is_bf16_supported(): + raise StarVLAError(f"CUDA device {device!r} does not support bfloat16") + + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + torch.use_deterministic_algorithms(True) + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + torch.backends.cudnn.benchmark = False + + +@contextlib.contextmanager +def _config_only_qwen_bootstrap(torch: Any, transformers: Any, qwen_dir: Path): + """Make the official wrapper construct Qwen topology without duplicate base weights. + + The subsequent strict load supplies every persistent parameter from the + SHA256-pinned StarVLA checkpoint. This is equivalent to StarVLA's released + loader after its bootstrap base weights are overwritten, while avoiding a + second 9 GB model download. + """ + + model_class = transformers.Qwen3VLForConditionalGeneration + had_local_override = "from_pretrained" in model_class.__dict__ + original_local_override = model_class.__dict__.get("from_pretrained") + + def from_config_only(model_id: str | os.PathLike[str], **kwargs: Any): + actual = Path(model_id).resolve() + if actual != qwen_dir.resolve(): + raise StarVLAError(f"official wrapper requested unexpected Qwen source: {actual}") + if kwargs.get("dtype") not in (None, torch.bfloat16): + raise StarVLAError(f"unexpected Qwen bootstrap dtype: {kwargs.get('dtype')!r}") + config = transformers.AutoConfig.from_pretrained( + actual, + local_files_only=True, + trust_remote_code=False, + ) + if getattr(config, "model_type", None) != "qwen3_vl": + raise StarVLAError(f"unexpected pinned Qwen model_type: {getattr(config, 'model_type', None)!r}") + previous_dtype = torch.get_default_dtype() + try: + torch.set_default_dtype(torch.bfloat16) + with transformers.modeling_utils.no_init_weights(): + model = model_class(config) + finally: + torch.set_default_dtype(previous_dtype) + return model + + model_class.from_pretrained = staticmethod(from_config_only) + try: + yield + finally: + if had_local_override: + model_class.from_pretrained = original_local_override + else: + delattr(model_class, "from_pretrained") + + +def _assert_module_origin(module: Any, source_dir: Path) -> None: + module_path = Path(module.__file__).resolve() + try: + module_path.relative_to(source_dir.resolve()) + except ValueError as exc: + raise StarVLAError(f"imported StarVLA module is outside the pinned checkout: {module_path}") from exc + + +@contextlib.contextmanager +def _official_qwen_model_alias(qwen_dir: Path): + """Expose the pinned local model under the case-sensitive name StarVLA dispatches on.""" + qwen_dir = qwen_dir.resolve() + with tempfile.TemporaryDirectory(prefix="starvla-qwen-alias-") as temporary: + alias = Path(temporary) / "Qwen3-VL-4B-Instruct" + alias.symlink_to(qwen_dir, target_is_directory=True) + if alias.resolve() != qwen_dir: + raise StarVLAError(f"temporary Qwen alias did not resolve to the pinned model: {alias}") + yield alias + + +def load_official_framework(paths: Mapping[str, Any], *, device: str) -> tuple[Any, dict[str, Any]]: + import torch + import transformers + + source_dir = Path(paths["source_dir"]) + if any(name == "starVLA" or name.startswith("starVLA.") for name in sys.modules): + raise StarVLAError("starVLA was imported before pinned-source verification") + sys.path.insert(0, str(source_dir)) + try: + from starVLA.model.framework import base_framework, share_tools + from starVLA.model.framework.VLM4A import QwenOFT + + _assert_module_origin(base_framework, source_dir) + _assert_module_origin(share_tools, source_dir) + _assert_module_origin(QwenOFT, source_dir) + config, norm_stats = share_tools.read_mode_config(str(paths["checkpoint"])) + qwen_dir = Path(paths["qwen_dir"]) + with _official_qwen_model_alias(qwen_dir) as qwen_alias: + config = base_framework.merge_config_overrides( + config, + [ + f"framework.qwenvl.base_vlm={qwen_alias}", + "framework.qwenvl.attn_implementation=sdpa", + ], + ) + expected_instruction = expected_framework_instruction( + config, + "contract probe", + int(config["framework"]["action_model"]["action_horizon"]), + ) + if ACTION_TOKEN * int(config["framework"]["action_model"]["action_horizon"]) not in expected_instruction: + raise StarVLAError("effective OFT config does not produce the pinned action-token contract") + + cfg = share_tools.dict_to_namespace(config) + cfg.trainer.pretrained_checkpoint = None + with _config_only_qwen_bootstrap(torch, transformers, qwen_dir): + framework = QwenOFT.Qwenvl_OFT(cfg) + + try: + state_dict = torch.load(paths["checkpoint"], map_location="cpu", weights_only=True) + except TypeError: + state_dict = torch.load(paths["checkpoint"], map_location="cpu") + if not isinstance(state_dict, Mapping) or not state_dict: + raise StarVLAError("official checkpoint did not contain a non-empty state_dict") + framework.load_state_dict(state_dict, strict=True) + del state_dict + gc.collect() + framework.norm_stats = norm_stats + + if type(framework).__name__ != "Qwenvl_OFT": + raise StarVLAError(f"unexpected official framework class: {type(framework).__name__}") + if int(framework.action_token_id) != ACTION_TOKEN_ID or framework.action_token != ACTION_TOKEN: + raise StarVLAError( + f"official action token mismatch: {framework.action_token!r}/{framework.action_token_id}" + ) + if int(framework.chunk_len) != 16: + raise StarVLAError(f"unexpected official OFT action horizon: {framework.chunk_len}") + qwen_dtypes = {parameter.dtype for parameter in framework.qwen_vl_interface.parameters()} + policy_dtypes = {parameter.dtype for parameter in framework.action_model.parameters()} + if qwen_dtypes != {torch.bfloat16}: + raise StarVLAError(f"unexpected Qwen parameter dtypes after strict load: {qwen_dtypes}") + if policy_dtypes != {torch.float32}: + raise StarVLAError(f"unexpected OFT parameter dtypes after strict load: {policy_dtypes}") + framework = framework.to(dtype=torch.bfloat16).to(device).eval() + if {parameter.dtype for parameter in framework.parameters()} != {torch.bfloat16}: + raise StarVLAError("official --use_bf16 cast did not cover the whole OFT model") + return framework, config + finally: + if sys.path and sys.path[0] == str(source_dir): + del sys.path[0] + + +def _tensor_to_array(tensor: Any) -> tuple[np.ndarray, str]: + source_dtype = str(tensor.dtype).removeprefix("torch.") + value = tensor.detach().cpu().contiguous() + if source_dtype == "bfloat16": + value = value.float() + return np.ascontiguousarray(value.numpy()), source_dtype + + +def _image_pixel_sha256(image: Any) -> str: + header = _canonical_json({"mode": image.mode, "size": list(image.size)}) + return _sha256_bytes(header + b"\x00" + image.tobytes()) + + +def _image_record(path: Path, image: Any) -> dict[str, Any]: + return { + "source_path": str(path.resolve()), + "source_size": path.stat().st_size, + "source_sha256": sha256_file(path), + "decoded_mode": image.mode, + "decoded_size": list(image.size), + "decoded_pixel_sha256": _image_pixel_sha256(image), + } + + +def _processed_image_records(images: Sequence[Any]) -> list[dict[str, Any]]: + return [ + { + "index": index, + "mode": image.mode, + "size": list(image.size), + "pixel_sha256": _image_pixel_sha256(image), + } + for index, image in enumerate(images) + ] + + +def run_official_forward(framework: Any, *, images: Sequence[Any], task: str) -> dict[str, Any]: + """Run Qwenvl_OFT.predict_action while capturing its real intermediate values.""" + + import torch + + captures: dict[str, Any] = {} + qwen = framework.qwen_vl_interface + action_model = framework.action_model + original_build = qwen.build_qwenvl_inputs + original_gather = framework._gather_action_token_embeddings + original_policy = action_model.predict_action + + def capture_build(*args: Any, **kwargs: Any): + batch_images = kwargs.get("images", args[0] if args else None) + instructions = kwargs.get("instructions", args[1] if len(args) > 1 else None) + captures["processed_images"] = list(batch_images[0]) + captures["framework_instructions"] = list(instructions) + result = original_build(*args, **kwargs) + captures["qwen_inputs"] = { + key: value.detach() + for key, value in result.items() + if isinstance(value, torch.Tensor) + } + return result + + def capture_gather(*args: Any, **kwargs: Any): + queries = original_gather(*args, **kwargs) + captures["action_queries_raw"] = queries.detach() + policy_dtype = next(action_model.parameters()).dtype + captures["policy_input_dtype"] = str(policy_dtype).removeprefix("torch.") + if queries.dtype != policy_dtype: + raise StarVLAError( + f"official whole-model BF16 dtype mismatch: queries={queries.dtype}, policy={policy_dtype}" + ) + return queries + + def capture_policy(*args: Any, **kwargs: Any): + captures["action_queries_policy"] = args[0].detach() + output = original_policy(*args, **kwargs) + captures["raw_policy"] = output.detach() + return output + + def capture_qwen_hidden(_module: Any, _inputs: Any, output: Any): + if not getattr(output, "hidden_states", None): + raise StarVLAError("official Qwen output did not include hidden_states") + captures["last_hidden_state"] = output.hidden_states[-1].detach() + + qwen.build_qwenvl_inputs = capture_build + framework._gather_action_token_embeddings = capture_gather + action_model.predict_action = capture_policy + hook = qwen.register_forward_hook(capture_qwen_hidden) + try: + result = framework.predict_action(examples=[{"image": list(images), "lang": task}]) + finally: + hook.remove() + qwen.build_qwenvl_inputs = original_build + framework._gather_action_token_embeddings = original_gather + action_model.predict_action = original_policy + + required = { + "processed_images", + "framework_instructions", + "qwen_inputs", + "action_queries_raw", + "action_queries_policy", + "raw_policy", + "last_hidden_state", + } + missing = sorted(required - set(captures)) + if missing: + raise StarVLAError(f"official OFT instrumentation did not capture: {missing}") + if "input_ids" not in captures["qwen_inputs"]: + raise StarVLAError("official Qwen preprocessing did not produce input_ids") + input_ids, _ = _tensor_to_array(captures["qwen_inputs"]["input_ids"]) + _, selected_positions = select_action_positions( + input_ids, + action_token_id=ACTION_TOKEN_ID, + chunk_len=int(framework.chunk_len), + ) + last_hidden = captures["last_hidden_state"] + positions = torch.as_tensor(selected_positions, device=last_hidden.device, dtype=torch.long) + expected_queries = last_hidden.gather( + 1, + positions.unsqueeze(-1).expand(-1, -1, last_hidden.shape[-1]), + ) + if not torch.equal(expected_queries, captures["action_queries_raw"]): + raise StarVLAError("captured action queries do not match final hidden state at selected token positions") + if captures["action_queries_raw"].dtype != torch.bfloat16: + raise StarVLAError(f"unexpected raw action-query dtype: {captures['action_queries_raw'].dtype}") + if captures["action_queries_policy"].dtype != torch.bfloat16: + raise StarVLAError(f"unexpected OFT input dtype: {captures['action_queries_policy'].dtype}") + if not torch.equal(captures["action_queries_raw"], captures["action_queries_policy"]): + raise StarVLAError("OFT policy input changed across the BF16 model boundary") + normalized = np.asarray(result.get("normalized_actions")) + raw_policy, _ = _tensor_to_array(captures["raw_policy"]) + expected_shape = (1, int(framework.chunk_len), int(action_model.action_dim)) + if normalized.shape != expected_shape: + raise StarVLAError(f"official OFT output shape mismatch: expected {expected_shape}, got {normalized.shape}") + if normalized.shape != raw_policy.shape or not np.array_equal(normalized, raw_policy): + raise StarVLAError("official normalized_actions differ from the raw OFT policy output") + if not np.isfinite(normalized).all(): + raise StarVLAError("official OFT policy produced NaN or infinite actions") + captures["normalized_actions"] = np.ascontiguousarray(normalized, dtype=np.float32) + return captures + + +def _render_model_prompt(framework: Any, processed_images: Sequence[Any], instruction: str) -> str: + messages = [ + { + "role": "user", + "content": [ + *({"type": "image", "image": image} for image in processed_images), + {"type": "text", "text": instruction}, + ], + } + ] + rendered = framework.qwen_vl_interface.processor.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True, + ) + if not isinstance(rendered, str): + raise StarVLAError(f"official processor returned a non-string rendered prompt: {type(rendered)}") + return rendered + + +def _build_arrays(captures: Mapping[str, Any], unnormalized: np.ndarray) -> tuple[dict[str, np.ndarray], dict[str, Any]]: + arrays: dict[str, np.ndarray] = {} + records: dict[str, Any] = {} + + def add(name: str, value: Any) -> None: + if isinstance(value, np.ndarray): + array = np.ascontiguousarray(value) + source_dtype = None + else: + array, source_dtype = _tensor_to_array(value) + arrays[name] = array + records[name] = _array_record(array, source_dtype=source_dtype) + + for key, tensor in sorted(captures["qwen_inputs"].items()): + add(f"qwen_input__{key}", tensor) + add("last_hidden_state", captures["last_hidden_state"]) + add("action_queries_raw", captures["action_queries_raw"]) + add("action_queries_policy", captures["action_queries_policy"]) + add("raw_policy", captures["raw_policy"]) + add("normalized_actions", captures["normalized_actions"]) + add("unnormalized_actions", np.ascontiguousarray(unnormalized)) + return arrays, records + + +def _runtime_record(torch: Any, transformers: Any, device: str) -> dict[str, Any]: + cuda_device = torch.device(device) + index = cuda_device.index if cuda_device.index is not None else torch.cuda.current_device() + properties = torch.cuda.get_device_properties(index) + return { + "python": platform.python_version(), + "platform": platform.platform(), + "torch": torch.__version__, + "torchvision": _distribution_version("torchvision"), + "transformers": transformers.__version__, + "numpy": np.__version__, + "pillow": _distribution_version("Pillow"), + "omegaconf": _distribution_version("omegaconf"), + "cuda_runtime": torch.version.cuda, + "cudnn": torch.backends.cudnn.version(), + "device": str(cuda_device), + "device_name": properties.name, + "compute_capability": [properties.major, properties.minor], + } + + +def _copy_inputs(staging: Path, image_paths: Sequence[Path]) -> list[str]: + inputs_dir = staging / "inputs" + inputs_dir.mkdir() + relative_paths = [] + for index, source in enumerate(image_paths): + suffix = source.suffix.lower() if source.suffix else ".img" + destination = inputs_dir / f"image-{index:02d}{suffix}" + shutil.copyfile(source, destination) + relative_paths.append(destination.relative_to(staging).as_posix()) + return relative_paths + + +def write_golden( + *, + output_dir: Path, + paths: Mapping[str, Any], + framework: Any, + config: Mapping[str, Any], + image_paths: Sequence[Path], + source_image_records: Sequence[Mapping[str, Any]], + task: str, + unnorm_key: str, + captures: Mapping[str, Any], + unnormalized: np.ndarray, +) -> Path: + import torch + import transformers + + output_dir = output_dir.resolve() + if output_dir.exists(): + raise StarVLAError(f"golden output directory already exists: {output_dir}") + output_dir.parent.mkdir(parents=True, exist_ok=True) + arrays, array_records = _build_arrays(captures, unnormalized) + input_ids = arrays["qwen_input__input_ids"] + all_positions, selected_positions = select_action_positions( + input_ids, + action_token_id=ACTION_TOKEN_ID, + chunk_len=int(framework.chunk_len), + ) + expected_instruction = expected_framework_instruction(config, task, int(framework.chunk_len)) + captured_instructions = captures["framework_instructions"] + if captured_instructions != [expected_instruction]: + raise StarVLAError( + f"official prompt contract changed: expected {expected_instruction!r}, got {captured_instructions!r}" + ) + model_instruction = expected_model_instruction(config, expected_instruction) + rendered_prompt = _render_model_prompt(framework, captures["processed_images"], model_instruction) + token_strings = framework.qwen_vl_interface.processor.tokenizer.convert_ids_to_tokens(input_ids[0].tolist()) + + identity = { + "schema_version": GOLDEN_SCHEMA_VERSION, + "variant": SUPPORTED_VARIANT, + "checkpoint_sha256": paths["variant"]["checkpoint"]["sha256"], + "starvla_revision": paths["catalog"]["source_revisions"]["starvla"], + "qwen_revision": paths["qwen"]["revision"], + "task": task, + "unnorm_key": unnorm_key, + "images": [record["source_sha256"] for record in source_image_records], + } + golden_id = _sha256_bytes(_canonical_json(identity)) + + with tempfile.TemporaryDirectory(prefix=f".{output_dir.name}.", dir=output_dir.parent) as temporary: + staging = Path(temporary) + copied_images = _copy_inputs(staging, image_paths) + tensor_path = staging / "tensors.npz" + np.savez(tensor_path, **arrays) + + image_records = [] + for index, record in enumerate(source_image_records): + copied = staging / copied_images[index] + image_records.append( + { + **record, + "artifact": copied_images[index], + "artifact_size": copied.stat().st_size, + "artifact_sha256": sha256_file(copied), + } + ) + + raw_policy = arrays["raw_policy"].tolist() + normalized = arrays["normalized_actions"].tolist() + unnormalized_list = arrays["unnormalized_actions"].tolist() + manifest = { + "schema_version": GOLDEN_SCHEMA_VERSION, + "kind": "starvla_oft_official_python_oracle", + "golden_id": golden_id, + "created_utc": dt.datetime.now(dt.timezone.utc).isoformat(), + "variant": SUPPORTED_VARIANT, + "model_type": paths["variant"]["model_type"], + "source": { + "catalog": str(paths["catalog_path"]), + "catalog_sha256": sha256_file(paths["catalog_path"]), + "bundle_uuid": official_bundle_uuid(paths["variant"], paths["catalog"]), + "starvla_repo_revision": paths["catalog"]["source_revisions"]["starvla"], + "starvla_checkout": str(paths["source_dir"]), + "checkpoint_repo_id": paths["variant"]["repo_id"], + "checkpoint_revision": paths["variant"]["revision"], + "checkpoint_path": str(paths["checkpoint"]), + "checkpoint_size": paths["variant"]["checkpoint"]["size"], + "checkpoint_sha256": paths["variant"]["checkpoint"]["sha256"], + "qwen_repo_id": paths["qwen"]["repo_id"], + "qwen_revision": paths["qwen"]["revision"], + "config_json_sha256": paths["variant"]["file_hashes"]["config.json"]["sha256"], + "config_yaml_sha256": paths["variant"]["file_hashes"]["config.yaml"]["sha256"], + "dataset_statistics_sha256": paths["variant"]["file_hashes"]["dataset_statistics.json"]["sha256"], + }, + "runtime": _runtime_record(torch, transformers, str(next(framework.parameters()).device)), + "determinism": { + "seed": 0, + "torch_deterministic_algorithms": True, + "cublas_workspace_config": os.environ.get("CUBLAS_WORKSPACE_CONFIG"), + "allow_tf32": False, + "attention_implementation": "sdpa", + }, + "compatibility": { + "qwen_bootstrap": ( + "config-only topology construction; every persistent parameter is then populated by " + "strict loading of the pinned official checkpoint" + ), + "whole_model_cast": { + "to": "bfloat16", + "reason": "official Bridge server launch uses --use_bf16", + }, + }, + "input": { + "task": task, + "unnorm_key": unnorm_key, + "images": image_records, + "processed_images": _processed_image_records(captures["processed_images"]), + }, + "model_contract": { + "framework_class": f"{type(framework).__module__}.{type(framework).__name__}", + "action_token": ACTION_TOKEN, + "action_token_id": ACTION_TOKEN_ID, + "action_horizon": int(framework.chunk_len), + "action_dim": int(framework.action_model.action_dim), + "qwen_hidden_dim": int(framework.qwen_vl_interface.model.config.hidden_size), + }, + "prompt": { + "framework_instruction": expected_instruction, + "model_instruction": model_instruction, + "rendered_chat_template": rendered_prompt, + }, + "tokens": { + "input_ids": input_ids.tolist(), + "token_strings": token_strings, + "all_action_token_positions": all_positions, + "selected_action_token_positions": selected_positions, + }, + "outputs": { + "raw_policy": raw_policy, + "normalized_actions": normalized, + "unnormalized_actions": unnormalized_list, + }, + "artifacts": { + "tensors": { + "path": tensor_path.name, + "size": tensor_path.stat().st_size, + "sha256": sha256_file(tensor_path), + "arrays": array_records, + } + }, + } + manifest_path = staging / "golden.json" + manifest_path.write_text( + json.dumps(manifest, ensure_ascii=False, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + Path(temporary).replace(output_dir) + return output_dir / "golden.json" + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Generate an auditable golden from the pinned official StarVLA Qwen3-VL OFT checkpoint." + ) + parser.add_argument("--image", action="append", default=[], type=Path, help="Ordered image input; repeat for views") + parser.add_argument("--task", help="Robot task instruction") + parser.add_argument("--unnorm-key", choices=("oxe_bridge", "oxe_rt1")) + parser.add_argument("--output-dir", type=Path) + parser.add_argument("--checkpoint-root", type=Path, default=Path("ckpts/starvla")) + parser.add_argument("--starvla-source", type=Path, default=None) + parser.add_argument("--device", default="cuda:0") + parser.add_argument( + "--preflight-only", + action="store_true", + help="Verify pinned source/assets/checkpoint/runtime without allocating the model", + ) + return parser + + +def _load_images(image_paths: Iterable[Path]) -> tuple[list[Any], list[dict[str, Any]]]: + try: + from PIL import Image + except ImportError as exc: + raise StarVLAError("Pillow is required to load oracle images") from exc + images = [] + records = [] + for path in image_paths: + path = path.resolve() + _ensure_regular_file(path, label="input image") + try: + with Image.open(path) as opened: + opened.load() + image = opened.copy() + except (OSError, ValueError) as exc: + raise StarVLAError(f"failed to decode input image {path}: {exc}") from exc + images.append(image) + records.append(_image_record(path, image)) + return images, records + + +def main(argv: Sequence[str] | None = None) -> int: + args = build_parser().parse_args(argv) + _require_isolated_python() + if not args.preflight_only: + missing = [ + name + for name, value in ( + ("--image", args.image), + ("--task", args.task), + ("--unnorm-key", args.unnorm_key), + ("--output-dir", args.output_dir), + ) + if not value + ] + if missing: + raise StarVLAError("golden generation requires " + ", ".join(missing)) + if not args.task.strip(): + raise StarVLAError("--task must not be empty") + checkpoint_root = Path(args.checkpoint_root).resolve() + output_dir = args.output_dir.resolve() + if output_dir == checkpoint_root or checkpoint_root in output_dir.parents: + raise StarVLAError("--output-dir must not be inside the pinned checkpoint source tree") + + checkpoint_root = args.checkpoint_root.resolve() + source_dir = args.starvla_source or checkpoint_root / "source" / "starvla" + paths = validate_official_inputs( + checkpoint_root=checkpoint_root, + source_dir=source_dir, + catalog_path=DEFAULT_CATALOG, + ) + + try: + import torch + import transformers + except ImportError as exc: + raise StarVLAError(f"official StarVLA runtime dependency is missing: {exc}") from exc + validate_runtime_versions( + torch_version=torch.__version__, + torchvision_version=_distribution_version("torchvision"), + transformers_version=transformers.__version__, + numpy_version=np.__version__, + ) + _configure_determinism(torch, seed=0, device=args.device) + if args.preflight_only: + print("Pinned StarVLA OFT oracle preflight passed.") + return 0 + + images, source_image_records = _load_images(args.image) + framework, config = load_official_framework(paths, device=args.device) + captures = run_official_forward(framework, images=images, task=args.task) + + source_dir = Path(paths["source_dir"]) + sys.path.insert(0, str(source_dir)) + try: + from deployment.model_server import policy_norm_processor + + _assert_module_origin(policy_norm_processor, source_dir) + normalizer = policy_norm_processor.PolicyNormProcessor( + str(paths["checkpoint"]), + unnorm_key=args.unnorm_key, + ) + normalized = captures["normalized_actions"] + unnormalized = normalizer.unapply_actions(normalized[0])[None, ...] + if unnormalized.shape != normalized.shape or not np.isfinite(unnormalized).all(): + raise StarVLAError( + f"official action unnormalization returned invalid values/shape: {unnormalized.shape}" + ) + finally: + if sys.path and sys.path[0] == str(source_dir): + del sys.path[0] + + manifest = write_golden( + output_dir=args.output_dir, + paths=paths, + framework=framework, + config=config, + image_paths=args.image, + source_image_records=source_image_records, + task=args.task, + unnorm_key=args.unnorm_key, + captures=captures, + unnormalized=np.ascontiguousarray(unnormalized), + ) + print(f"Wrote official StarVLA OFT golden: {manifest}") + return 0 + + +if __name__ == "__main__": + try: + raise SystemExit(main()) + except StarVLAError as exc: + raise SystemExit(f"error: {exc}") from exc diff --git a/tools/hf2gguf/starvla/serve_starvla_oft_reference.py b/tools/hf2gguf/starvla/serve_starvla_oft_reference.py new file mode 100644 index 0000000..99f0787 --- /dev/null +++ b/tools/hf2gguf/starvla/serve_starvla_oft_reference.py @@ -0,0 +1,1152 @@ +#!/usr/bin/env python3 +"""Serve a pinned official StarVLA OFT checkpoint over robot.cpp protocol v3. + +This process is the Python-reference backend for closed-loop parity evaluation. +It deliberately shares the robot.cpp client protocol and SimplerEnv adapter, so +the only policy variable is the original PyTorch checkpoint versus GGUF. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import logging +import math +import os +import socket +import struct +import subprocess +import sys +import tempfile +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable, Mapping, Sequence + +import numpy as np + + +TOOLS_DIR = Path(__file__).resolve().parent +REPO_ROOT = TOOLS_DIR.parents[2] +for search_path in (TOOLS_DIR, REPO_ROOT): + if str(search_path) not in sys.path: + sys.path.insert(0, str(search_path)) + +from generate_starvla_oft_golden import ( # noqa: E402 + _assert_module_origin, + _configure_determinism, + _distribution_version, + _require_isolated_python, + load_official_framework as load_qwen3_official_framework, + validate_official_inputs as validate_qwen3_official_inputs, + validate_runtime_versions, +) +from generate_starvla_qwen25_oft_golden import ( # noqa: E402 + EXPECTED_QWEN_VL_UTILS_VERSION, + LEGACY_UNNORM_PROFILES, + legacy_normalization_contract, + load_official_framework as load_qwen25_official_framework, + unnormalize_legacy_actions, + validate_local_inputs as validate_qwen25_local_inputs, +) +from robot_client.python import model_client as wire # noqa: E402 +from starvla_checkpoint import ( # noqa: E402 + DEFAULT_CATALOG, + StarVLAError, + get_qwen_asset, + get_variant, + load_catalog, + official_bundle_uuid, + sha256_file, + verify_catalog_files, +) + + +SERVER_METADATA_SCHEMA_VERSION = 2 +MODEL_TYPE = "starvla" +FRAMEWORK = "oft" +REFERENCE_BACKEND = "local-python-checkpoint-reference" +REFERENCE_PURPOSE = "bridge-only" +DEFAULT_IMAGE_NAME = "image_0" +DEFAULT_UNNORM_KEY = "oxe_bridge" +QWEN25_DEFAULT_UNNORM_KEY = "bridge_dataset" +SUPPORTED_VARIANTS = ("oft", "qwen25_oft") + +STATUS_BAD_REQUEST = 1 +STATUS_BAD_VERSION = 3 +STATUS_PAYLOAD_TOO_BIG = 4 +STATUS_INTERNAL_ERROR = 5 +MAX_PAYLOAD_BYTES = 256 * 1024 * 1024 + + +class ProtocolError(ValueError): + """A malformed or contract-invalid protocol request.""" + + +class PayloadTooBig(ProtocolError): + pass + + +@dataclass(frozen=True) +class RequestHeader: + magic: int + version: int + header_size: int + op: int + flags: int + request_id: int + status: int + payload_len: int + reserved: int + + +@dataclass(frozen=True) +class WireImage: + name: str + width: int + height: int + channels: int + stride_bytes: int + data: bytes + + def to_rgb_array(self) -> np.ndarray: + rows = np.frombuffer( + self.data, + dtype=np.uint8, + count=self.stride_bytes * self.height, + ).reshape(self.height, self.stride_bytes) + packed = rows[:, : self.width * self.channels] + return np.ascontiguousarray(packed.reshape(self.height, self.width, self.channels)) + + +@dataclass(frozen=True) +class PredictRequest: + images: tuple[WireImage, ...] + state: tuple[float, ...] + task: str + + +@dataclass(frozen=True) +class PredictResult: + actions: np.ndarray + metrics: Mapping[str, float] + + +def _canonical_json_bytes(value: Any) -> bytes: + return json.dumps( + value, + ensure_ascii=True, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + + +def _canonical_sha256(value: Any) -> str: + return hashlib.sha256(_canonical_json_bytes(value)).hexdigest() + + +def _decode_utf8(value: bytes, label: str) -> str: + try: + return value.decode("utf-8", errors="strict") + except UnicodeDecodeError as exc: + raise ProtocolError(f"{label} is not valid UTF-8") from exc + + +def _asset_manifest(entry: Mapping[str, Any]) -> dict[str, Any]: + return { + "repo_id": entry["repo_id"], + "revision": entry["revision"], + "files": { + relative: { + "size": int(entry["file_hashes"][relative]["size"]), + "sha256": str(entry["file_hashes"][relative]["sha256"]), + } + for relative in sorted(entry.get("files", [])) + }, + } + + +def _git_tree_sha1(source_dir: Path) -> str: + try: + result = subprocess.run( + ["git", "-C", str(source_dir), "rev-parse", "HEAD^{tree}"], + check=True, + capture_output=True, + text=True, + ) + except (OSError, subprocess.CalledProcessError) as exc: + raise StarVLAError(f"failed to resolve pinned StarVLA Git tree: {exc}") from exc + value = result.stdout.strip() + if len(value) != 40 or any(ch not in "0123456789abcdef" for ch in value): + raise StarVLAError(f"invalid pinned StarVLA Git tree SHA1: {value!r}") + return value + + +def _git_tracked_index_sha256(source_dir: Path) -> str: + """Hash the clean checkout's modes, paths, and Git blob identities.""" + + try: + result = subprocess.run( + ["git", "-C", str(source_dir), "ls-files", "-s", "-z"], + check=True, + capture_output=True, + ) + except (OSError, subprocess.CalledProcessError) as exc: + raise StarVLAError( + f"failed to hash pinned StarVLA tracked-file manifest: {exc}" + ) from exc + if not result.stdout: + raise StarVLAError("pinned StarVLA tracked-file manifest is empty") + return hashlib.sha256(result.stdout).hexdigest() + + +def default_unnorm_key_for_variant(variant: str) -> str: + if variant == "oft": + return DEFAULT_UNNORM_KEY + if variant == "qwen25_oft": + return QWEN25_DEFAULT_UNNORM_KEY + raise StarVLAError( + f"unsupported OFT reference variant {variant!r}; expected one of {SUPPORTED_VARIANTS}" + ) + + +def validate_reference_inputs( + *, + checkpoint_root: Path, + source_dir: Path, + variant_name: str, + catalog_path: Path = DEFAULT_CATALOG, +) -> dict[str, Any]: + """Resolve and verify one official OFT checkpoint without accepting aliases.""" + + if variant_name == "oft": + return validate_qwen3_official_inputs( + checkpoint_root=checkpoint_root, + source_dir=source_dir, + catalog_path=catalog_path, + ) + if variant_name != "qwen25_oft": + raise StarVLAError( + f"unsupported OFT reference variant {variant_name!r}; " + f"expected one of {SUPPORTED_VARIANTS}" + ) + + catalog = load_catalog(catalog_path) + variant = get_variant(catalog, variant_name) + qwen_asset_name, qwen = get_qwen_asset(catalog, variant) + if ( + variant.get("framework") != FRAMEWORK + or variant.get("model_type") != MODEL_TYPE + or variant.get("backbone") != "qwen2_5_vl" + or qwen_asset_name != "qwen2_5_vl_3b_instruct" + ): + raise StarVLAError("catalog Qwen2.5 OFT identity is incompatible") + + checkpoint_root = checkpoint_root.resolve() + expected_source = (checkpoint_root / "source" / "starvla").resolve() + source_dir = source_dir.resolve() + if source_dir != expected_source: + raise StarVLAError( + f"StarVLA source must be the canonical checkout {expected_source}, got {source_dir}" + ) + + policy_dir = ( + checkpoint_root / "sources" / variant["directory"] / variant["revision"] + ) + qwen_dir = checkpoint_root / "sources" / qwen["directory"] / qwen["revision"] + checkpoint = policy_dir / variant["checkpoint"]["path"] + verify_catalog_files(policy_dir, variant) + verify_catalog_files(qwen_dir, qwen) + local = validate_qwen25_local_inputs( + checkpoint=checkpoint, + qwen_model=qwen_dir, + source_dir=source_dir, + expected_checkpoint_sha256=str(variant["checkpoint"]["sha256"]), + expected_checkpoint_size=int(variant["checkpoint"]["size"]), + expected_source_revision=str(catalog["source_revisions"]["starvla"]), + ) + return { + **local, + "catalog": catalog, + "variant": variant, + "qwen": qwen, + "policy_dir": policy_dir, + "catalog_path": catalog_path.resolve(), + } + + +def build_preflight_record(paths: Mapping[str, Any]) -> dict[str, Any]: + variant = paths["variant"] + qwen = paths["qwen"] + return { + "schema_version": 1, + "ready": True, + "variant": variant["_catalog_key"], + "model_type": variant["model_type"], + "framework": variant["framework"], + "backbone": variant.get("backbone", "qwen3_vl"), + "checkpoint": { + "repo_id": variant["repo_id"], + "revision": variant["revision"], + "path": str(Path(paths["checkpoint"]).resolve()), + "size": int(variant["checkpoint"]["size"]), + "sha256": variant["checkpoint"]["sha256"], + }, + "qwen": { + "repo_id": qwen["repo_id"], + "revision": qwen["revision"], + "path": str(Path(paths["qwen_dir"]).resolve()), + }, + "starvla": { + "revision": paths["catalog"]["source_revisions"]["starvla"], + "path": str(Path(paths["source_dir"]).resolve()), + }, + "catalog": { + "path": str(Path(paths["catalog_path"]).resolve()), + "sha256": sha256_file(Path(paths["catalog_path"])), + }, + } + + +def apply_official_bf16(framework: Any, torch: Any) -> Any: + """Match the official Bridge server's ``--use_bf16`` whole-model cast.""" + + framework = framework.to(dtype=torch.bfloat16) + dtypes = {parameter.dtype for parameter in framework.parameters()} + if dtypes != {torch.bfloat16}: + raise StarVLAError(f"official OFT model must be entirely bfloat16, got {dtypes}") + return framework + + +def build_runtime_metadata( + torch: Any, + transformers: Any, + *, + device: str, +) -> dict[str, Any]: + """Record the exact local Python runtime used by the Bridge reference.""" + + torch_device = torch.device(device) + cudnn_version = torch.backends.cudnn.version() + record: dict[str, Any] = { + "python_full_version": sys.version, + "torch": str(torch.__version__), + "torch_cuda": ( + None if torch.version.cuda is None else str(torch.version.cuda) + ), + "cudnn": None if cudnn_version is None else int(cudnn_version), + "transformers": str(transformers.__version__), + "pillow": _distribution_version("Pillow"), + "numpy": str(np.__version__), + "device": str(torch_device), + "gpu_name": None, + "compute_capability": None, + } + if torch_device.type == "cuda": + index = ( + torch_device.index + if torch_device.index is not None + else torch.cuda.current_device() + ) + properties = torch.cuda.get_device_properties(index) + record["gpu_name"] = str(properties.name) + record["compute_capability"] = [ + int(properties.major), + int(properties.minor), + ] + return record + + +def build_server_metadata( + paths: Mapping[str, Any], + framework: Any, + *, + default_unnorm_key: str, + source_tree_sha1: str, + source_tracked_index_sha256: str, + runtime: Mapping[str, Any], +) -> dict[str, Any]: + catalog = paths["catalog"] + variant = paths["variant"] + qwen = paths["qwen"] + variant_name = str(variant["_catalog_key"]) + backbone = str(variant.get("backbone", "qwen3_vl")) + profiles = [str(value) for value in framework.norm_stats.keys()] + if not profiles or len(set(profiles)) != len(profiles): + raise StarVLAError(f"invalid official normalization profiles: {profiles}") + expected_profiles = ( + list(LEGACY_UNNORM_PROFILES) + if variant_name == "qwen25_oft" + else ["oxe_bridge", "oxe_rt1"] + ) + if profiles != expected_profiles: + raise StarVLAError( + f"unexpected {variant_name} normalization profiles: " + f"expected {expected_profiles}, got {profiles}" + ) + if default_unnorm_key not in profiles: + raise StarVLAError( + f"default unnorm key {default_unnorm_key!r} is not in {profiles}" + ) + + qwen_dtypes = sorted( + {str(parameter.dtype).removeprefix("torch.") for parameter in framework.qwen_vl_interface.parameters()} + ) + policy_dtypes = sorted( + {str(parameter.dtype).removeprefix("torch.") for parameter in framework.action_model.parameters()} + ) + if qwen_dtypes != ["bfloat16"] or policy_dtypes != ["bfloat16"]: + raise StarVLAError( + f"unexpected loaded dtype profile: qwen={qwen_dtypes}, oft={policy_dtypes}" + ) + chunk_size = int(framework.chunk_len) + action_dim = int(framework.action_model.action_dim) + if chunk_size <= 0 or action_dim <= 0: + raise StarVLAError( + f"invalid official action contract: chunk={chunk_size}, dim={action_dim}" + ) + + checkpoint_sha256 = str(variant["checkpoint"]["sha256"]) + checkpoint_revision = str(variant["revision"]) + qwen_manifest = _asset_manifest(qwen) + policy_manifest = _asset_manifest(variant) + starvla_revision = str(catalog["source_revisions"]["starvla"]) + model_info = { + "model_type": MODEL_TYPE, + "framework": FRAMEWORK, + "bundle_uuid": official_bundle_uuid(variant, catalog), + "checkpoint_sha256": checkpoint_sha256, + "checkpoint_revision": checkpoint_revision, + "qwen_revision": str(qwen["revision"]), + "starvla_revision": starvla_revision, + "image_names": [DEFAULT_IMAGE_NAME], + "state_supported": False, + "state_dimension_dynamic": False, + "state_dim": 0, + "chunk_size": chunk_size, + "action_dim": action_dim, + "normalization_profiles": profiles, + "default_unnorm_key": default_unnorm_key, + } + return { + "schema_version": SERVER_METADATA_SCHEMA_VERSION, + "protocol_version": wire.VERSION, + "backend": REFERENCE_BACKEND, + "purpose": REFERENCE_PURPOSE, + "catalog_variant": variant_name, + "backbone": backbone, + "runtime": dict(runtime), + "model_info": model_info, + "checkpoint": { + "repo_id": variant["repo_id"], + "revision": checkpoint_revision, + "path": str(Path(paths["checkpoint"]).resolve()), + "size": int(variant["checkpoint"]["size"]), + "sha256": checkpoint_sha256, + "asset_manifest_sha256": _canonical_sha256(policy_manifest), + }, + "qwen": { + "repo_id": qwen["repo_id"], + "revision": qwen["revision"], + "bootstrap_assets_manifest_sha256": _canonical_sha256(qwen_manifest), + "bootstrap_assets": qwen_manifest["files"], + }, + "starvla_source": { + "revision": starvla_revision, + "commit_sha": starvla_revision, + "git_tree_sha1": source_tree_sha1, + "tracked_index_manifest_sha256": source_tracked_index_sha256, + "path": str(Path(paths["source_dir"]).resolve()), + }, + "catalog": { + "path": str(Path(paths["catalog_path"]).resolve()), + "sha256": sha256_file(Path(paths["catalog_path"])), + }, + "dtype_profile": { + "qwen_parameters": "bfloat16", + "qwen_action_queries": "bfloat16", + "oft_input_cast": None, + "oft_parameters": "bfloat16", + "wire_actions": "float32", + "whole_model_cast": True, + }, + "action_contract": { + "chunk_size": chunk_size, + "action_dim": action_dim, + }, + "normalization": { + "implementation": ( + "released_q01_q99_masked_with_binary_unmasked_dimensions" + if variant_name == "qwen25_oft" + else "official PolicyNormProcessor" + ), + "available_unnorm_keys": profiles, + "default_unnorm_key": default_unnorm_key, + "runtime_robot_profile_aliases": ( + dict(LEGACY_UNNORM_PROFILES) + if variant_name == "qwen25_oft" + else {profile: profile for profile in profiles} + ), + }, + } + + +def decode_request_header(raw: bytes) -> RequestHeader: + if len(raw) != wire.HEADER_SIZE: + raise ProtocolError("short header") + return RequestHeader(*wire.HEADER.unpack(raw)) + + +def validate_request_header(header: RequestHeader) -> None: + if header.magic != wire.MAGIC: + raise ProtocolError("bad magic") + if header.version != wire.VERSION: + raise ProtocolError("bad protocol version") + if header.header_size != wire.HEADER_SIZE: + raise ProtocolError("bad header size") + if header.flags != 0 or header.status != wire.STATUS_OK or header.reserved != 0: + raise ProtocolError("request header flags/status/reserved must be zero") + if header.payload_len > MAX_PAYLOAD_BYTES: + raise PayloadTooBig("payload too large") + + +def decode_predict_request(payload: bytes) -> PredictRequest: + if len(payload) < wire.PREDICT_REQ_V2_FIXED.size: + raise ProtocolError("short predict request") + image_count, state_count, task_len = wire.PREDICT_REQ_V2_FIXED.unpack_from(payload) + if image_count == 0: + raise ProtocolError("predict request requires at least one image") + + offset = wire.PREDICT_REQ_V2_FIXED.size + remaining = len(payload) - offset + if image_count > remaining // wire.PREDICT_REQ_V2_IMAGE.size: + raise ProtocolError("image count exceeds predict request metadata") + + metadata: list[tuple[int, int, int, int, int, int]] = [] + for index in range(image_count): + ( + image_format, + name_len, + width, + height, + channels, + stride_bytes, + data_len, + ) = wire.PREDICT_REQ_V2_IMAGE.unpack_from(payload, offset) + offset += wire.PREDICT_REQ_V2_IMAGE.size + if image_format != wire.IMAGE_RAW_RGB_U8: + raise ProtocolError(f"image[{index}] has an unsupported image format") + if width <= 0 or height <= 0 or channels != 3: + raise ProtocolError(f"image[{index}] has invalid raw RGB dimensions") + packed_stride = width * channels + if stride_bytes < packed_stride: + raise ProtocolError(f"image[{index}] has an invalid stride_bytes") + if data_len < stride_bytes * height: + raise ProtocolError( + f"image[{index}] data is smaller than stride_bytes * height" + ) + metadata.append((name_len, width, height, channels, stride_bytes, data_len)) + + body_size = ( + state_count * 4 + + task_len + + sum(name_len + data_len for name_len, *_rest, data_len in metadata) + ) + if body_size != len(payload) - offset: + raise ProtocolError("predict request fields do not exactly match payload") + + state: tuple[float, ...] + if state_count: + state = tuple(struct.unpack_from(f"<{state_count}f", payload, offset)) + else: + state = () + offset += state_count * 4 + if any(not math.isfinite(value) for value in state): + raise ProtocolError("state contains a non-finite value") + + task = _decode_utf8(payload[offset : offset + task_len], "task") + offset += task_len + + images: list[WireImage] = [] + for index, (name_len, width, height, channels, stride_bytes, data_len) in enumerate(metadata): + name = _decode_utf8(payload[offset : offset + name_len], f"image[{index}] name") + offset += name_len + data = bytes(payload[offset : offset + data_len]) + offset += data_len + images.append( + WireImage( + name=name, + width=width, + height=height, + channels=channels, + stride_bytes=stride_bytes, + data=data, + ) + ) + if offset != len(payload): + raise ProtocolError("trailing bytes in predict request") + return PredictRequest( + images=tuple(images), + state=state, + task=task, + ) + + +def encode_predict_response(result: PredictResult) -> bytes: + actions = np.asarray(result.actions, dtype=np.float32) + if actions.ndim != 2 or actions.shape[0] <= 0 or actions.shape[1] <= 0: + raise ProtocolError(f"invalid action matrix: {actions.shape}") + if not np.isfinite(actions).all(): + raise ProtocolError("actions contain non-finite values") + chunk_size, action_dim = (int(actions.shape[0]), int(actions.shape[1])) + action_count = actions.size + metric_rows: list[tuple[bytes, float]] = [] + for name in sorted(result.metrics): + name_bytes = str(name).encode("utf-8") + if not name_bytes: + raise ProtocolError("metric name is empty") + value = float(result.metrics[name]) + if not math.isfinite(value): + raise ProtocolError(f"metric {name!r} is non-finite") + metric_rows.append((name_bytes, value)) + + output = bytearray( + wire.PREDICT_RESP_FIXED.pack( + chunk_size, + action_dim, + action_count, + len(metric_rows), + ) + ) + for name, value in metric_rows: + output += wire.PREDICT_RESP_METRIC.pack(len(name), value) + output += name + output += np.ascontiguousarray(actions, dtype=" None: + if (processor_factory is None) == (action_unnormalizer is None): + raise StarVLAError( + "reference policy requires exactly one action unnormalization implementation" + ) + self.framework = framework + self.action_unnormalizer = action_unnormalizer + self.metadata = dict(metadata) + self.model_info = dict(self.metadata["model_info"]) + self.unnorm_key = str(self.model_info["default_unnorm_key"]) + self.processor = None + if processor_factory is not None: + self.processor = processor_factory( + str(checkpoint), unnorm_key=self.unnorm_key + ) + if getattr(self.processor, "unnorm_key", self.unnorm_key) != self.unnorm_key: + raise StarVLAError("PolicyNormProcessor selected the wrong profile") + + def reset(self) -> None: + # OFT is stateless; this method intentionally preserves loaded weights. + return None + + def predict(self, request: PredictRequest) -> PredictResult: + if len(request.images) != 1: + raise ProtocolError( + f"OFT Bridge reference requires exactly one image, got {len(request.images)}" + ) + image = request.images[0] + if image.name != DEFAULT_IMAGE_NAME: + raise ProtocolError( + f"OFT Bridge reference requires image name {DEFAULT_IMAGE_NAME!r}, got {image.name!r}" + ) + if request.state: + raise ProtocolError("OFT Bridge reference does not accept robot state") + if not request.task.strip(): + raise ProtocolError("task must not be empty") + + try: + from PIL import Image + except ImportError as exc: + raise RuntimeError("Pillow is required for OFT reference inference") from exc + + pil_image = Image.fromarray(image.to_rgb_array(), mode="RGB") + total_started = time.perf_counter() + forward_started = time.perf_counter() + output = self.framework.predict_action( + examples=[{"image": [pil_image], "lang": request.task}] + ) + forward_ms = (time.perf_counter() - forward_started) * 1000.0 + if not isinstance(output, Mapping) or "normalized_actions" not in output: + raise RuntimeError("official OFT forward did not return normalized_actions") + normalized = np.asarray(output["normalized_actions"]) + expected_shape = ( + 1, + int(self.model_info["chunk_size"]), + int(self.model_info["action_dim"]), + ) + if normalized.shape != expected_shape or not np.isfinite(normalized).all(): + raise RuntimeError( + f"official OFT returned invalid normalized actions: {normalized.shape}" + ) + + unnorm_started = time.perf_counter() + if self.action_unnormalizer is None: + assert self.processor is not None + actions = np.asarray( + self.processor.unapply_actions(normalized[0]), + dtype=np.float32, + ) + else: + actions = np.asarray( + self.action_unnormalizer(normalized, self.unnorm_key), + dtype=np.float32, + ) + unnorm_ms = (time.perf_counter() - unnorm_started) * 1000.0 + if actions.shape != expected_shape[1:] or not np.isfinite(actions).all(): + raise RuntimeError( + f"official PolicyNormProcessor returned invalid actions: {actions.shape}" + ) + total_ms = (time.perf_counter() - total_started) * 1000.0 + return PredictResult( + actions=np.ascontiguousarray(actions), + metrics={ + "python_forward_ms": forward_ms, + "python_unnorm_ms": unnorm_ms, + "model_total_ms": total_ms, + }, + ) + + +class ProtocolApplication: + def __init__(self, policy: Any): + self.policy = policy + self.shutdown_requested = False + + def dispatch( + self, + op: int, + payload: bytes, + *, + server_recv_ms: float = 0.0, + ) -> tuple[bytes, bool]: + if op == wire.OP_HEALTH: + if payload: + raise ProtocolError("health request payload must be empty") + return f"ok policy={self.policy.model_info['model_type']}".encode(), False + if op == wire.OP_RESET: + if payload: + raise ProtocolError("reset request payload must be empty") + self.policy.reset() + return b"ok", False + if op == wire.OP_SHUTDOWN: + if payload: + raise ProtocolError("shutdown request payload must be empty") + self.shutdown_requested = True + return b"ok", True + if op == wire.OP_PREDICT: + request = decode_predict_request(payload) + predict_started = time.perf_counter() + result = self.policy.predict(request) + server_predict_ms = (time.perf_counter() - predict_started) * 1000.0 + metrics = dict(result.metrics) + metrics.update( + { + "server_queue_ms": 0.0, + "server_predict_ms": server_predict_ms, + "server_recv_ms": float(server_recv_ms), + } + ) + return ( + encode_predict_response( + PredictResult( + actions=result.actions, + metrics=metrics, + ) + ), + False, + ) + raise ProtocolError("unknown op") + + +def _recv_exact(sock: socket.socket, length: int, *, allow_initial_eof: bool = False) -> bytes | None: + chunks: list[bytes] = [] + remaining = length + while remaining: + chunk = sock.recv(remaining) + if not chunk: + if allow_initial_eof and remaining == length: + return None + raise ConnectionError("peer closed connection") + chunks.append(chunk) + remaining -= len(chunk) + return b"".join(chunks) + + +def _response_header( + request: RequestHeader, + *, + status: int, + payload_len: int, +) -> bytes: + return wire.HEADER.pack( + wire.MAGIC, + wire.VERSION, + wire.HEADER_SIZE, + request.op, + 0, + request.request_id, + status, + payload_len, + 0, + ) + + +class ReferenceProtocolServer: + """Small sequential TCP server matching robot_server/session.cpp semantics.""" + + def __init__( + self, + policy: Any, + *, + host: str = "127.0.0.1", + port: int = 5555, + backlog: int = 16, + ) -> None: + if host != "127.0.0.1": + raise ValueError("Python reference server only listens on 127.0.0.1") + if port < 0 or port > 65535: + raise ValueError("port must be in 0..65535") + self.application = ProtocolApplication(policy) + self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + self.socket.bind((host, port)) + self.socket.listen(backlog) + self.address = self.socket.getsockname() + self._closed = False + + def close(self) -> None: + if not self._closed: + self.socket.close() + self._closed = True + + def _send( + self, + client: socket.socket, + request: RequestHeader, + status: int, + payload: bytes, + ) -> None: + client.sendall( + _response_header(request, status=status, payload_len=len(payload)) + ) + if payload: + client.sendall(payload) + + def _handle_client(self, client: socket.socket) -> None: + while not self.application.shutdown_requested: + recv_started = time.perf_counter() + raw_header = _recv_exact( + client, + wire.HEADER_SIZE, + allow_initial_eof=True, + ) + if raw_header is None: + return + request = decode_request_header(raw_header) + if request.magic != wire.MAGIC: + return + if request.version != wire.VERSION: + self._send( + client, + request, + STATUS_BAD_VERSION, + b"bad protocol version", + ) + return + try: + validate_request_header(request) + except PayloadTooBig as exc: + self._send(client, request, STATUS_PAYLOAD_TOO_BIG, str(exc).encode("utf-8")) + return + except ProtocolError as exc: + self._send(client, request, STATUS_BAD_REQUEST, str(exc).encode("utf-8")) + return + payload = _recv_exact(client, request.payload_len) if request.payload_len else b"" + assert payload is not None + server_recv_ms = (time.perf_counter() - recv_started) * 1000.0 + try: + response, should_shutdown = self.application.dispatch( + request.op, + payload, + server_recv_ms=server_recv_ms, + ) + status = wire.STATUS_OK + except ProtocolError as exc: + response = str(exc).encode("utf-8") + should_shutdown = False + status = STATUS_BAD_REQUEST + except Exception as exc: + logging.exception("Python reference inference failed") + response = f"Python reference inference failed: {exc}".encode("utf-8") + should_shutdown = False + status = STATUS_INTERNAL_ERROR + self._send(client, request, status, response) + if should_shutdown: + return + + def serve_forever(self) -> None: + try: + while not self.application.shutdown_requested: + client, peer = self.socket.accept() + logging.debug("protocol connection from %s:%s", *peer) + with client: + try: + self._handle_client(client) + except (ConnectionError, OSError): + logging.debug("protocol peer disconnected", exc_info=True) + finally: + self.close() + + +def load_pinned_reference_policy( + *, + checkpoint_root: Path, + starvla_source: Path | None, + device: str, + noise_seed: int, + default_unnorm_key: str, + variant_name: str = "oft", +) -> PinnedOFTReferencePolicy: + source_dir = starvla_source or checkpoint_root / "source" / "starvla" + paths = validate_reference_inputs( + checkpoint_root=checkpoint_root, + source_dir=Path(source_dir), + variant_name=variant_name, + catalog_path=DEFAULT_CATALOG, + ) + try: + import torch + import transformers + except ImportError as exc: + raise StarVLAError(f"official StarVLA runtime dependency is missing: {exc}") from exc + validate_runtime_versions( + torch_version=torch.__version__, + torchvision_version=_distribution_version("torchvision"), + transformers_version=transformers.__version__, + numpy_version=np.__version__, + ) + if ( + variant_name == "qwen25_oft" + and _distribution_version("qwen-vl-utils") + != EXPECTED_QWEN_VL_UTILS_VERSION + ): + raise StarVLAError( + "official Qwen2.5 OFT reference requires qwen-vl-utils " + f"{EXPECTED_QWEN_VL_UTILS_VERSION}, got " + f"{_distribution_version('qwen-vl-utils')}" + ) + _configure_determinism(torch, seed=noise_seed, device=device) + framework_loader = ( + load_qwen25_official_framework + if variant_name == "qwen25_oft" + else load_qwen3_official_framework + ) + framework, _config = framework_loader(paths, device=device) + framework = apply_official_bf16(framework, torch) + + source_dir = Path(paths["source_dir"]) + processor_factory: Callable[..., Any] | None = None + action_unnormalizer: Callable[[np.ndarray, str], np.ndarray] | None = None + if variant_name == "qwen25_oft": + for profile in framework.norm_stats: + legacy_normalization_contract(framework.norm_stats, str(profile)) + + def qwen25_unnormalizer( + normalized: np.ndarray, profile: str + ) -> np.ndarray: + return unnormalize_legacy_actions( + normalized, + framework.norm_stats, + profile, + )[0] + + action_unnormalizer = qwen25_unnormalizer + else: + sys.path.insert(0, str(source_dir)) + try: + from deployment.model_server import policy_norm_processor + + _assert_module_origin(policy_norm_processor, source_dir) + processor_factory = policy_norm_processor.PolicyNormProcessor + finally: + if sys.path and sys.path[0] == str(source_dir): + del sys.path[0] + + metadata = build_server_metadata( + paths, + framework, + default_unnorm_key=default_unnorm_key, + source_tree_sha1=_git_tree_sha1(source_dir), + source_tracked_index_sha256=_git_tracked_index_sha256(source_dir), + runtime=build_runtime_metadata( + torch, + transformers, + device=device, + ), + ) + return PinnedOFTReferencePolicy( + framework=framework, + processor_factory=processor_factory, + checkpoint=Path(paths["checkpoint"]), + metadata=metadata, + action_unnormalizer=action_unnormalizer, + ) + + +def write_metadata(path: Path, metadata: Mapping[str, Any]) -> None: + path = path.resolve() + path.parent.mkdir(parents=True, exist_ok=True) + serialized = json.dumps(metadata, indent=2, sort_keys=True) + "\n" + with tempfile.NamedTemporaryFile( + mode="w", + encoding="utf-8", + dir=path.parent, + prefix=f".{path.name}.", + delete=False, + ) as handle: + temporary = Path(handle.name) + handle.write(serialized) + handle.flush() + os.fsync(handle.fileno()) + try: + os.replace(temporary, path) + finally: + temporary.unlink(missing_ok=True) + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description=( + "Serve a pinned official StarVLA Qwen-VL OFT checkpoint over " + "robot.cpp protocol v4." + ) + ) + parser.add_argument( + "--variant", + choices=SUPPORTED_VARIANTS, + default="oft", + help="Catalog variant; qwen25_oft uses the plain Qwen2.5-VL assets.", + ) + parser.add_argument("--checkpoint-root", type=Path, default=Path("ckpts/starvla")) + parser.add_argument("--starvla-source", type=Path) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, default=5555) + parser.add_argument( + "--unnorm-key", + help="Defaults to oxe_bridge for oft and bridge_dataset for qwen25_oft.", + ) + parser.add_argument( + "--preflight", + action="store_true", + help="Verify pinned local inputs and print provenance without loading the model.", + ) + parser.add_argument( + "--noise-seed", + type=int, + default=0, + help="Deterministic runtime seed; OFT inference itself has no sampled noise.", + ) + parser.add_argument( + "--metadata-output", + type=Path, + help="Optional atomic JSON record of all pinned identities and dtype contracts.", + ) + parser.add_argument("--verbosity", type=int, default=0) + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + args = build_parser().parse_args(argv) + _require_isolated_python() + if args.host != "127.0.0.1": + raise StarVLAError("--host must be 127.0.0.1") + if args.port <= 0 or args.port > 65535: + raise StarVLAError("--port must be in 1..65535") + if args.verbosity < 0: + raise StarVLAError("--verbosity must be non-negative") + logging.basicConfig( + level=logging.DEBUG if args.verbosity else logging.INFO, + format="%(asctime)s %(levelname)s %(message)s", + force=True, + ) + default_unnorm_key = ( + args.unnorm_key or default_unnorm_key_for_variant(args.variant) + ) + + checkpoint_root = args.checkpoint_root.resolve() + source_dir = ( + args.starvla_source.resolve() + if args.starvla_source + else checkpoint_root / "source" / "starvla" + ) + if args.preflight: + paths = validate_reference_inputs( + checkpoint_root=checkpoint_root, + source_dir=source_dir, + variant_name=args.variant, + catalog_path=DEFAULT_CATALOG, + ) + record = build_preflight_record(paths) + serialized = json.dumps(record, indent=2, sort_keys=True) + "\n" + if args.metadata_output is not None: + write_metadata(args.metadata_output, record) + sys.stdout.write(serialized) + return 0 + + policy = load_pinned_reference_policy( + checkpoint_root=checkpoint_root, + starvla_source=source_dir, + device=args.device, + noise_seed=args.noise_seed, + default_unnorm_key=default_unnorm_key, + variant_name=args.variant, + ) + if args.metadata_output is not None: + write_metadata(args.metadata_output, policy.metadata) + logging.info( + "loaded pinned OFT Python reference metadata=%s", + _canonical_json_bytes(policy.metadata).decode("ascii"), + ) + + server = ReferenceProtocolServer(policy, host=args.host, port=args.port) + logging.info( + "Python reference server listening on %s:%d model=%s variant=%s", + server.address[0], + server.address[1], + MODEL_TYPE, + args.variant, + ) + try: + server.serve_forever() + except KeyboardInterrupt: + logging.info("Python reference server interrupted") + server.close() + return 0 + + +if __name__ == "__main__": + try: + raise SystemExit(main()) + except StarVLAError as exc: + raise SystemExit(f"error: {exc}") from exc From 190d5beac1301b72b566baa19c930b689dcb9269 Mon Sep 17 00:00:00 2001 From: JJJYmmm <1650675829@qq.com> Date: Mon, 10 Aug 2026 12:42:56 +0800 Subject: [PATCH 03/11] starvla: add Qwen2.5-VL OFT policy --- .../generate_starvla_qwen25_oft_golden.py | 986 ++++++++++++++++++ 1 file changed, 986 insertions(+) create mode 100644 tools/hf2gguf/starvla/generate_starvla_qwen25_oft_golden.py diff --git a/tools/hf2gguf/starvla/generate_starvla_qwen25_oft_golden.py b/tools/hf2gguf/starvla/generate_starvla_qwen25_oft_golden.py new file mode 100644 index 0000000..87dfdfa --- /dev/null +++ b/tools/hf2gguf/starvla/generate_starvla_qwen25_oft_golden.py @@ -0,0 +1,986 @@ +#!/usr/bin/env python3 +"""Generate a local-Python oracle from the official Qwen2.5-VL OFT .pt file. + +Unlike the Qwen3-VL catalog-driven oracle, this entry point binds the golden +directly to an explicitly supplied local checkpoint and processor directory. +The full StarVLA checkpoint supplies all model parameters; the Qwen directory +is used only for the model topology and processor/tokenizer assets. +""" + +from __future__ import annotations + +import argparse +import contextlib +import datetime as dt +import gc +import json +import os +import shutil +import sys +import tempfile +from pathlib import Path +from typing import Any, Iterable, Mapping, Sequence + +import numpy as np + + +TOOLS_DIR = Path(__file__).resolve().parent +if str(TOOLS_DIR) not in sys.path: + sys.path.insert(0, str(TOOLS_DIR)) + +from generate_starvla_oft_golden import ( # noqa: E402 + _array_record, + _assert_module_origin, + _canonical_json, + _configure_determinism, + _distribution_version, + _ensure_regular_file, + _image_pixel_sha256, + _require_isolated_python, + _runtime_record, + _sha256_bytes, + _tensor_to_array, + expected_framework_instruction, + expected_model_instruction, + select_action_positions, + validate_runtime_versions, +) +from starvla_checkpoint import StarVLAError, sha256_file # noqa: E402 + + +SCHEMA_VERSION = 1 +GOLDEN_KIND = "starvla_qwen25_oft_local_pt_python_oracle" +MODEL_TYPE = "starvla" +BACKBONE = "qwen2_5_vl" +ACTION_TOKEN = chr(0x1F50D) +ACTION_RELATIVE_L2_LIMIT = 0.03 +EXPECTED_QWEN_VL_UTILS_VERSION = "0.0.14" +LEGACY_UNNORM_PROFILES = { + "bridge_dataset": "oxe_bridge", + "fractal20220817_data": "oxe_rt1", +} + +OFFICIAL_CHECKPOINT_REPO_ID = "StarVLA/Qwen-OFT-Bridge-RT-1" +OFFICIAL_CHECKPOINT_REVISION = "11fa6440835ba3e912de43cfe8521043360ffc02" +OFFICIAL_CHECKPOINT_FILENAME = "steps_10000_pytorch_model.pt" +OFFICIAL_CHECKPOINT_SIZE = 8_215_912_766 +OFFICIAL_CHECKPOINT_SHA256 = "51fe8d22c8d57116c2f59c5fdb24323fa3411149e888b807edba99b8354e0861" +OFFICIAL_QWEN_REVISION = "66285546d2b821cf421d4f5eb2576359d3770cd3" +OFFICIAL_BUNDLE_UUID = "90d105ae-00fa-580c-8751-9f931e324c3b" + +QWEN_PROCESSOR_REQUIRED = { + "config.json", + "preprocessor_config.json", + "tokenizer_config.json", +} +QWEN_PROCESSOR_OPTIONAL = { + "added_tokens.json", + "chat_template.json", + "merges.txt", + "special_tokens_map.json", + "tokenizer.json", + "vocab.json", +} + + +def _run_git(source_dir: Path, *arguments: str) -> str: + import subprocess + + try: + result = subprocess.run( + ["git", "-C", str(source_dir), *arguments], + check=True, + capture_output=True, + text=True, + ) + except (OSError, subprocess.CalledProcessError) as exc: + raise StarVLAError(f"failed to inspect StarVLA checkout {source_dir}: {exc}") from exc + return result.stdout.strip() + + +def _verify_clean_source(source_dir: Path, expected_revision: str | None) -> str: + source_dir = source_dir.resolve() + if not (source_dir / ".git").is_dir(): + raise StarVLAError(f"StarVLA source is not a Git checkout: {source_dir}") + revision = _run_git(source_dir, "rev-parse", "HEAD") + if expected_revision is not None and revision != expected_revision: + raise StarVLAError( + f"StarVLA source revision mismatch: expected {expected_revision}, got {revision}" + ) + changes = _run_git(source_dir, "status", "--porcelain=v1", "--untracked-files=all") + if changes: + raise StarVLAError(f"StarVLA source checkout is not clean:\n{changes}") + return revision + + +def _load_json_object(path: Path, *, label: str) -> dict[str, Any]: + _ensure_regular_file(path, label=label) + try: + value = json.loads(path.read_text(encoding="utf-8")) + except (OSError, UnicodeError, json.JSONDecodeError) as exc: + raise StarVLAError(f"failed to parse {label} {path}: {exc}") from exc + if not isinstance(value, dict): + raise StarVLAError(f"{label} root must be an object") + return value + + +def _qwen_asset_records(qwen_dir: Path) -> tuple[list[dict[str, Any]], str]: + qwen_dir = qwen_dir.resolve() + missing = sorted(name for name in QWEN_PROCESSOR_REQUIRED if not (qwen_dir / name).is_file()) + if missing: + raise StarVLAError("Qwen2.5 processor directory is missing: " + ", ".join(missing)) + + tokenizer_files = {"tokenizer.json", "vocab.json"} + if not any((qwen_dir / name).is_file() for name in tokenizer_files): + raise StarVLAError("Qwen2.5 processor directory needs tokenizer.json or vocab.json") + + records: list[dict[str, Any]] = [] + for name in sorted(QWEN_PROCESSOR_REQUIRED | QWEN_PROCESSOR_OPTIONAL): + path = qwen_dir / name + if not path.exists(): + continue + _ensure_regular_file(path, label=f"Qwen2.5 asset {name}") + records.append( + { + "path": name, + "size": path.stat().st_size, + "sha256": sha256_file(path), + } + ) + digest = _sha256_bytes(_canonical_json(records)) + return records, digest + + +def validate_local_inputs( + *, + checkpoint: Path, + qwen_model: Path, + source_dir: Path, + expected_checkpoint_sha256: str = OFFICIAL_CHECKPOINT_SHA256, + expected_checkpoint_size: int = OFFICIAL_CHECKPOINT_SIZE, + expected_source_revision: str | None = None, +) -> dict[str, Any]: + checkpoint = checkpoint.resolve() + _ensure_regular_file(checkpoint, label="Qwen2.5 OFT checkpoint") + if checkpoint.suffix != ".pt": + raise StarVLAError(f"Qwen2.5 OFT reference checkpoint must be a .pt file: {checkpoint}") + if Path(f"{checkpoint}.aria2").exists(): + raise StarVLAError(f"checkpoint download is incomplete: {checkpoint}.aria2 exists") + actual_size = checkpoint.stat().st_size + if actual_size != expected_checkpoint_size: + raise StarVLAError( + f"checkpoint size mismatch: expected {expected_checkpoint_size}, got {actual_size}" + ) + actual_sha256 = sha256_file(checkpoint) + if actual_sha256 != expected_checkpoint_sha256: + raise StarVLAError( + f"checkpoint SHA256 mismatch: expected {expected_checkpoint_sha256}, got {actual_sha256}" + ) + if checkpoint.parent.name != "checkpoints" or len(checkpoint.parents) < 2: + raise StarVLAError( + "checkpoint must use the StarVLA layout /checkpoints/.pt" + ) + + run_dir = checkpoint.parents[1] + config_yaml = run_dir / "config.yaml" + dataset_statistics = run_dir / "dataset_statistics.json" + _ensure_regular_file(config_yaml, label="checkpoint config.yaml") + stats = _load_json_object(dataset_statistics, label="checkpoint dataset_statistics.json") + if not stats: + raise StarVLAError("checkpoint dataset_statistics.json must not be empty") + + qwen_model = qwen_model.resolve() + if not qwen_model.is_dir(): + raise StarVLAError(f"Qwen2.5 processor/model directory does not exist: {qwen_model}") + qwen_config_path = qwen_model / "config.json" + qwen_config = _load_json_object(qwen_config_path, label="Qwen2.5 config.json") + if qwen_config.get("model_type") != BACKBONE: + raise StarVLAError( + f"Qwen config model_type must be {BACKBONE!r}, got {qwen_config.get('model_type')!r}" + ) + hidden_size = qwen_config.get("hidden_size") + if hidden_size != 2048: + raise StarVLAError(f"official Qwen2.5-VL 3B hidden_size must be 2048, got {hidden_size!r}") + qwen_assets, qwen_assets_sha256 = _qwen_asset_records(qwen_model) + + source_dir = source_dir.resolve() + source_revision = _verify_clean_source(source_dir, expected_source_revision) + return { + "checkpoint": checkpoint, + "checkpoint_size": actual_size, + "checkpoint_sha256": actual_sha256, + "run_dir": run_dir, + "config_yaml": config_yaml, + "dataset_statistics": dataset_statistics, + "norm_stats": stats, + "qwen_dir": qwen_model, + "qwen_config": qwen_config, + "qwen_assets": qwen_assets, + "qwen_assets_sha256": qwen_assets_sha256, + "source_dir": source_dir, + "source_revision": source_revision, + } + + +def legacy_normalization_contract( + norm_stats: Mapping[str, Any], unnorm_key: str +) -> dict[str, Any]: + if unnorm_key not in LEGACY_UNNORM_PROFILES: + raise StarVLAError( + f"Qwen2.5 OFT unnorm_key must be one of {sorted(LEGACY_UNNORM_PROFILES)}, " + f"got {unnorm_key!r}" + ) + value = norm_stats.get(unnorm_key) + if not isinstance(value, Mapping) or not isinstance(value.get("action"), Mapping): + raise StarVLAError(f"dataset statistics has no {unnorm_key}.action object") + action = value["action"] + try: + q01 = np.asarray(action["q01"], dtype=np.float32) + q99 = np.asarray(action["q99"], dtype=np.float32) + mask = np.asarray(action["mask"], dtype=np.bool_) + except (KeyError, TypeError, ValueError) as exc: + raise StarVLAError(f"invalid legacy action statistics for {unnorm_key}: {exc}") from exc + if q01.shape != (7,) or q99.shape != (7,) or mask.shape != (7,): + raise StarVLAError( + f"legacy action statistics must be 7D, got {q01.shape}/{q99.shape}/{mask.shape}" + ) + if not np.isfinite(q01).all() or not np.isfinite(q99).all(): + raise StarVLAError("legacy Qwen2.5 OFT q01/q99 statistics must be finite") + if np.any(q99[mask] <= q01[mask]): + raise StarVLAError("legacy Qwen2.5 OFT masked q99 values must exceed q01") + return { + "stats_key": unnorm_key, + "runtime_robot_profile": LEGACY_UNNORM_PROFILES[unnorm_key], + "method": "q01_q99_masked_with_binary_unmasked_dimensions", + "binary_threshold": 0.5, + "q01": q01.tolist(), + "q99": q99.tolist(), + "mask": mask.tolist(), + } + + +def unnormalize_legacy_actions( + normalized: np.ndarray, + norm_stats: Mapping[str, Any], + unnorm_key: str, +) -> np.ndarray: + """Mirror the released checkpoint's q99 + binary action transform.""" + + contract = legacy_normalization_contract(norm_stats, unnorm_key) + actions = np.ascontiguousarray(normalized, dtype=np.float32) + if actions.ndim != 3 or actions.shape[0] != 1 or actions.shape[-1] != 7: + raise StarVLAError( + f"normalized Qwen2.5 OFT actions must have shape [1,T,7], got {actions.shape}" + ) + q01 = np.asarray(contract["q01"], dtype=np.float32) + q99 = np.asarray(contract["q99"], dtype=np.float32) + mask = np.asarray(contract["mask"], dtype=np.bool_) + result = np.empty_like(actions) + result[..., mask] = ( + (actions[..., mask] + np.float32(1.0)) + / np.float32(2.0) + * (q99[mask] - q01[mask]) + + q01[mask] + ) + result[..., ~mask] = ( + actions[..., ~mask] > np.float32(contract["binary_threshold"]) + ).astype(np.float32) + if not np.isfinite(result).all(): + raise StarVLAError("legacy Qwen2.5 OFT unnormalization produced non-finite actions") + return np.ascontiguousarray(result) + + +@contextlib.contextmanager +def _config_only_qwen25_bootstrap(torch: Any, transformers: Any, qwen_dir: Path): + """Build Qwen2.5 topology without loading a duplicate base weight set.""" + + model_class = transformers.Qwen2_5_VLForConditionalGeneration + had_override = "from_pretrained" in model_class.__dict__ + original_override = model_class.__dict__.get("from_pretrained") + + def from_config_only(model_id: str | os.PathLike[str], **kwargs: Any): + actual = Path(model_id).resolve() + if actual != qwen_dir.resolve(): + raise StarVLAError(f"official wrapper requested unexpected Qwen source: {actual}") + torch_dtype = kwargs.get("torch_dtype") + if torch_dtype not in (None, "auto", torch.bfloat16): + raise StarVLAError(f"unexpected Qwen bootstrap torch_dtype: {torch_dtype!r}") + config = transformers.AutoConfig.from_pretrained( + actual, + local_files_only=True, + trust_remote_code=False, + ) + declared_model_type = getattr(type(config), "model_type", None) + runtime_model_type = getattr(config, "model_type", None) + text_config = getattr(config, "text_config", config) + vision_config = getattr(config, "vision_config", None) + config_contract = { + "declared_model_type": declared_model_type, + "runtime_model_type": runtime_model_type, + "hidden_size": getattr(text_config, "hidden_size", None), + "layer_count": getattr(text_config, "num_hidden_layers", None), + "vocab_size": getattr(text_config, "vocab_size", None), + "vision_hidden_size": getattr(vision_config, "hidden_size", 1280), + "vision_depth": getattr(vision_config, "depth", 32), + "vision_output_size": getattr(vision_config, "out_hidden_size", 2048), + } + expected_contract = { + "declared_model_type": BACKBONE, + # Transformers 4.57 delegates this instance property to text_config. + "runtime_model_type": runtime_model_type, + "hidden_size": 2048, + "layer_count": 36, + "vocab_size": 151936, + "vision_hidden_size": 1280, + "vision_depth": 32, + "vision_output_size": 2048, + } + if ( + runtime_model_type not in {BACKBONE, "qwen2_5_vl_text"} + or config_contract != expected_contract + ): + raise StarVLAError( + f"unexpected local Qwen config contract: {config_contract}" + ) + previous_dtype = torch.get_default_dtype() + try: + torch.set_default_dtype(torch.bfloat16) + with transformers.modeling_utils.no_init_weights(): + model = model_class(config) + finally: + torch.set_default_dtype(previous_dtype) + return model + + model_class.from_pretrained = staticmethod(from_config_only) + try: + yield + finally: + if had_override: + model_class.from_pretrained = original_override + else: + delattr(model_class, "from_pretrained") + + +@contextlib.contextmanager +def _official_qwen25_alias(qwen_dir: Path): + """Give the local processor directory the dispatch name used by StarVLA.""" + + qwen_dir = qwen_dir.resolve() + with tempfile.TemporaryDirectory(prefix="starvla-qwen25-alias-") as temporary: + alias = Path(temporary) / "Qwen2.5-VL-3B-Instruct" + alias.symlink_to(qwen_dir, target_is_directory=True) + if alias.resolve() != qwen_dir: + raise StarVLAError(f"temporary Qwen2.5 alias has the wrong target: {alias}") + yield alias + + +def load_official_framework(paths: Mapping[str, Any], *, device: str) -> tuple[Any, dict[str, Any]]: + import torch + import transformers + + source_dir = Path(paths["source_dir"]) + if any(name == "starVLA" or name.startswith("starVLA.") for name in sys.modules): + raise StarVLAError("starVLA was imported before local source verification") + sys.path.insert(0, str(source_dir)) + try: + from starVLA.model.framework import base_framework, share_tools + from starVLA.model.framework.VLM4A import QwenOFT + + _assert_module_origin(base_framework, source_dir) + _assert_module_origin(share_tools, source_dir) + _assert_module_origin(QwenOFT, source_dir) + config, norm_stats = share_tools.read_mode_config(str(paths["checkpoint"])) + with _official_qwen25_alias(Path(paths["qwen_dir"])) as qwen_alias: + config = base_framework.merge_config_overrides( + config, + [ + f"framework.qwenvl.base_vlm={qwen_alias}", + "framework.qwenvl.attn_implementation=sdpa", + ], + ) + cfg = share_tools.dict_to_namespace(config) + cfg.trainer.pretrained_checkpoint = None + with _config_only_qwen25_bootstrap(torch, transformers, Path(paths["qwen_dir"])): + framework = QwenOFT.Qwenvl_OFT(cfg) + + try: + state_dict = torch.load(paths["checkpoint"], map_location="cpu", weights_only=True) + except TypeError: + state_dict = torch.load(paths["checkpoint"], map_location="cpu") + if not isinstance(state_dict, Mapping) or not state_dict: + raise StarVLAError("official checkpoint did not contain a non-empty state_dict") + framework.load_state_dict(state_dict, strict=True) + del state_dict + gc.collect() + framework.norm_stats = norm_stats + + if type(framework).__name__ != "Qwenvl_OFT": + raise StarVLAError(f"unexpected official framework class: {type(framework).__name__}") + if framework.action_token != ACTION_TOKEN: + raise StarVLAError(f"unexpected OFT action token: {framework.action_token!r}") + token_ids = framework.qwen_vl_interface.processor.tokenizer( + ACTION_TOKEN, add_special_tokens=False + )["input_ids"] + if token_ids != [int(framework.action_token_id)]: + raise StarVLAError(f"Qwen2.5 action token is not one tokenizer token: {token_ids}") + if int(framework.chunk_len) != 16: + raise StarVLAError(f"official OFT action horizon must be 16, got {framework.chunk_len}") + if int(framework.action_model.action_dim) != 7: + raise StarVLAError( + f"official OFT action dimension must be 7, got {framework.action_model.action_dim}" + ) + hidden_size = int(framework.qwen_vl_interface.model.config.hidden_size) + if hidden_size != 2048: + raise StarVLAError(f"official Qwen2.5 hidden size must be 2048, got {hidden_size}") + qwen_dtypes = {parameter.dtype for parameter in framework.qwen_vl_interface.parameters()} + policy_dtypes = {parameter.dtype for parameter in framework.action_model.parameters()} + if qwen_dtypes != {torch.bfloat16}: + raise StarVLAError(f"unexpected Qwen2.5 parameter dtypes: {qwen_dtypes}") + if policy_dtypes != {torch.float32}: + raise StarVLAError(f"unexpected OFT parameter dtypes: {policy_dtypes}") + framework = framework.to(dtype=torch.bfloat16).to(device).eval() + if {parameter.dtype for parameter in framework.parameters()} != {torch.bfloat16}: + raise StarVLAError("official --use_bf16 cast did not cover the whole OFT model") + return framework, config + finally: + if sys.path and sys.path[0] == str(source_dir): + del sys.path[0] + + +def run_official_forward(framework: Any, *, images: Sequence[Any], task: str) -> dict[str, Any]: + """Execute Qwenvl_OFT.predict_action and capture the exact policy boundary.""" + + import torch + + captures: dict[str, Any] = {} + qwen = framework.qwen_vl_interface + action_model = framework.action_model + original_build = qwen.build_qwenvl_inputs + original_gather = framework._gather_action_token_embeddings + original_policy = action_model.predict_action + + def capture_build(*args: Any, **kwargs: Any): + batch_images = kwargs.get("images", args[0] if args else None) + instructions = kwargs.get("instructions", args[1] if len(args) > 1 else None) + captures["processed_images"] = list(batch_images[0]) + captures["framework_instructions"] = list(instructions) + result = original_build(*args, **kwargs) + captures["qwen_inputs"] = { + key: value.detach() + for key, value in result.items() + if isinstance(value, torch.Tensor) + } + return result + + def capture_gather(*args: Any, **kwargs: Any): + queries = original_gather(*args, **kwargs) + captures["action_queries_raw"] = queries.detach() + policy_dtype = next(action_model.parameters()).dtype + captures["policy_input_dtype"] = str(policy_dtype).removeprefix("torch.") + if queries.dtype != policy_dtype: + raise StarVLAError( + f"official whole-model BF16 dtype mismatch: queries={queries.dtype}, policy={policy_dtype}" + ) + return queries + + def capture_policy(*args: Any, **kwargs: Any): + captures["action_queries_policy"] = args[0].detach() + output = original_policy(*args, **kwargs) + captures["raw_policy"] = output.detach() + return output + + def capture_hidden(_module: Any, _inputs: Any, output: Any): + if not getattr(output, "hidden_states", None): + raise StarVLAError("official Qwen2.5 output did not include hidden_states") + captures["last_hidden_state"] = output.hidden_states[-1].detach() + + qwen.build_qwenvl_inputs = capture_build + framework._gather_action_token_embeddings = capture_gather + action_model.predict_action = capture_policy + hook = qwen.register_forward_hook(capture_hidden) + try: + result = framework.predict_action(examples=[{"image": list(images), "lang": task}]) + finally: + hook.remove() + qwen.build_qwenvl_inputs = original_build + framework._gather_action_token_embeddings = original_gather + action_model.predict_action = original_policy + + required = { + "processed_images", + "framework_instructions", + "qwen_inputs", + "action_queries_raw", + "action_queries_policy", + "raw_policy", + "last_hidden_state", + } + missing = sorted(required - set(captures)) + if missing: + raise StarVLAError(f"official Qwen2.5 OFT instrumentation missed: {missing}") + + input_ids, _ = _tensor_to_array(captures["qwen_inputs"]["input_ids"]) + _, selected = select_action_positions( + input_ids, + action_token_id=int(framework.action_token_id), + chunk_len=int(framework.chunk_len), + ) + last_hidden = captures["last_hidden_state"] + positions = torch.as_tensor(selected, device=last_hidden.device, dtype=torch.long) + expected_queries = last_hidden.gather( + 1, positions.unsqueeze(-1).expand(-1, -1, last_hidden.shape[-1]) + ) + if not torch.equal(expected_queries, captures["action_queries_raw"]): + raise StarVLAError("captured action queries do not match Qwen2.5 result_norm positions") + if captures["action_queries_raw"].dtype != torch.bfloat16: + raise StarVLAError( + f"unexpected raw action-query dtype: {captures['action_queries_raw'].dtype}" + ) + if captures["action_queries_policy"].dtype != torch.bfloat16: + raise StarVLAError( + f"unexpected OFT policy input dtype: {captures['action_queries_policy'].dtype}" + ) + if not torch.equal(captures["action_queries_raw"], captures["action_queries_policy"]): + raise StarVLAError("OFT policy input changed across the BF16 model boundary") + + normalized = np.asarray(result.get("normalized_actions")) + raw_policy, _ = _tensor_to_array(captures["raw_policy"]) + expected_shape = (1, int(framework.chunk_len), int(action_model.action_dim)) + if normalized.shape != expected_shape: + raise StarVLAError( + f"official OFT output shape mismatch: expected {expected_shape}, got {normalized.shape}" + ) + if normalized.shape != raw_policy.shape or not np.array_equal(normalized, raw_policy): + raise StarVLAError("normalized_actions differ from the captured OFT policy output") + if not np.isfinite(normalized).all(): + raise StarVLAError("official Qwen2.5 OFT produced non-finite actions") + captures["normalized_actions"] = np.ascontiguousarray(normalized, dtype=np.float32) + return captures + + +def _load_images(image_paths: Iterable[Path]) -> tuple[list[Any], list[dict[str, Any]]]: + try: + from PIL import Image + except ImportError as exc: + raise StarVLAError("Pillow is required to load oracle images") from exc + + images: list[Any] = [] + records: list[dict[str, Any]] = [] + for path in image_paths: + path = path.resolve() + _ensure_regular_file(path, label="input image") + try: + with Image.open(path) as opened: + opened.load() + image = opened.copy() + except (OSError, ValueError) as exc: + raise StarVLAError(f"failed to decode input image {path}: {exc}") from exc + images.append(image) + records.append( + { + "source_path": str(path), + "source_size": path.stat().st_size, + "source_sha256": sha256_file(path), + "decoded_mode": image.mode, + "decoded_size": list(image.size), + "decoded_pixel_sha256": _image_pixel_sha256(image), + } + ) + return images, records + + +def _render_model_prompt(framework: Any, images: Sequence[Any], instruction: str) -> str: + messages = [ + { + "role": "user", + "content": [ + *({"type": "image", "image": image} for image in images), + {"type": "text", "text": instruction}, + ], + } + ] + rendered = framework.qwen_vl_interface.processor.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True + ) + if not isinstance(rendered, str): + raise StarVLAError(f"Qwen2.5 processor returned a non-string prompt: {type(rendered)}") + return rendered + + +def _build_arrays( + captures: Mapping[str, Any], unnormalized: np.ndarray +) -> tuple[dict[str, np.ndarray], dict[str, Any]]: + arrays: dict[str, np.ndarray] = {} + records: dict[str, Any] = {} + + def add(name: str, value: Any) -> None: + if isinstance(value, np.ndarray): + array = np.ascontiguousarray(value) + source_dtype = None + else: + array, source_dtype = _tensor_to_array(value) + arrays[name] = array + records[name] = _array_record(array, source_dtype=source_dtype) + + for key, tensor in sorted(captures["qwen_inputs"].items()): + add(f"qwen_input__{key}", tensor) + add("last_hidden_state", captures["last_hidden_state"]) + add("action_queries_raw", captures["action_queries_raw"]) + add("action_queries_policy", captures["action_queries_policy"]) + add("raw_policy", captures["raw_policy"]) + add("normalized_actions", captures["normalized_actions"]) + add("unnormalized_actions", np.ascontiguousarray(unnormalized, dtype=np.float32)) + return arrays, records + + +def write_golden( + *, + output_dir: Path, + paths: Mapping[str, Any], + framework: Any, + config: Mapping[str, Any], + image_paths: Sequence[Path], + source_image_records: Sequence[Mapping[str, Any]], + task: str, + unnorm_key: str, + captures: Mapping[str, Any], + unnormalized: np.ndarray, +) -> Path: + import torch + import transformers + + output_dir = output_dir.resolve() + if output_dir.exists(): + raise StarVLAError(f"golden output directory already exists: {output_dir}") + output_dir.parent.mkdir(parents=True, exist_ok=True) + + arrays, array_records = _build_arrays(captures, unnormalized) + input_ids = arrays["qwen_input__input_ids"] + if len(image_paths) != 1 or len(source_image_records) != 1: + raise StarVLAError("Qwen2.5 OFT Bridge parity requires exactly one image") + image_grid = arrays.get("qwen_input__image_grid_thw") + pixel_values = arrays.get("qwen_input__pixel_values") + if ( + image_grid is None + or image_grid.shape != (1, 3) + or not np.issubdtype(image_grid.dtype, np.integer) + ): + raise StarVLAError( + "official Qwen2.5 processor must return one image_grid_thw row" + ) + grid_thw = [int(value) for value in image_grid[0]] + if any(value <= 0 for value in grid_thw) or grid_thw[0] != 1: + raise StarVLAError(f"unexpected Qwen2.5 image_grid_thw: {grid_thw}") + patch_count = int(np.prod(grid_thw, dtype=np.int64)) + if ( + pixel_values is None + or pixel_values.ndim != 2 + or pixel_values.shape[0] != patch_count + ): + raise StarVLAError( + "Qwen2.5 pixel_values do not match the processor image grid" + ) + patch_size = 14 + spatial_merge_size = 2 + if patch_count % (spatial_merge_size * spatial_merge_size) != 0: + raise StarVLAError("Qwen2.5 image grid is not divisible by spatial merge size") + resized_size = [grid_thw[2] * patch_size, grid_thw[1] * patch_size] + merged_image_token_count = patch_count // ( + spatial_merge_size * spatial_merge_size + ) + action_token_id = int(framework.action_token_id) + all_positions, selected_positions = select_action_positions( + input_ids, + action_token_id=action_token_id, + chunk_len=int(framework.chunk_len), + ) + framework_instruction = expected_framework_instruction( + config, task, int(framework.chunk_len) + ) + if captures["framework_instructions"] != [framework_instruction]: + raise StarVLAError("official Qwen2.5 OFT prompt construction drifted") + model_instruction = expected_model_instruction(config, framework_instruction) + rendered_prompt = _render_model_prompt( + framework, captures["processed_images"], model_instruction + ) + token_strings = framework.qwen_vl_interface.processor.tokenizer.convert_ids_to_tokens( + input_ids[0].tolist() + ) + + identity = { + "schema_version": SCHEMA_VERSION, + "kind": GOLDEN_KIND, + "checkpoint_sha256": paths["checkpoint_sha256"], + "starvla_revision": paths["source_revision"], + "qwen_assets_sha256": paths["qwen_assets_sha256"], + "task": task, + "unnorm_key": unnorm_key, + "state": [], + "images": [record["source_sha256"] for record in source_image_records], + } + golden_id = _sha256_bytes(_canonical_json(identity)) + + with tempfile.TemporaryDirectory( + prefix=f".{output_dir.name}.", dir=output_dir.parent + ) as temporary: + staging = Path(temporary) + inputs_dir = staging / "inputs" + inputs_dir.mkdir() + image_records: list[dict[str, Any]] = [] + for index, (source_path, source_record) in enumerate( + zip(image_paths, source_image_records, strict=True) + ): + suffix = source_path.suffix.lower() if source_path.suffix else ".img" + artifact = inputs_dir / f"image-{index:02d}{suffix}" + shutil.copyfile(source_path, artifact) + image_records.append( + { + **source_record, + "artifact": artifact.relative_to(staging).as_posix(), + "artifact_size": artifact.stat().st_size, + "artifact_sha256": sha256_file(artifact), + } + ) + + tensors_path = staging / "tensors.npz" + np.savez(tensors_path, **arrays) + manifest = { + "schema_version": SCHEMA_VERSION, + "kind": GOLDEN_KIND, + "golden_id": golden_id, + "created_utc": dt.datetime.now(dt.timezone.utc).isoformat(), + "model_type": MODEL_TYPE, + "backbone": BACKBONE, + "source": { + "checkpoint_repo_id": OFFICIAL_CHECKPOINT_REPO_ID, + "checkpoint_revision": OFFICIAL_CHECKPOINT_REVISION, + "checkpoint_filename": OFFICIAL_CHECKPOINT_FILENAME, + "checkpoint_path": str(paths["checkpoint"]), + "checkpoint_size": paths["checkpoint_size"], + "checkpoint_sha256": paths["checkpoint_sha256"], + "config_yaml_path": str(paths["config_yaml"]), + "config_yaml_size": paths["config_yaml"].stat().st_size, + "config_yaml_sha256": sha256_file(paths["config_yaml"]), + "dataset_statistics_path": str(paths["dataset_statistics"]), + "dataset_statistics_size": paths["dataset_statistics"].stat().st_size, + "dataset_statistics_sha256": sha256_file(paths["dataset_statistics"]), + "qwen_model_path": str(paths["qwen_dir"]), + "qwen_assets": paths["qwen_assets"], + "qwen_assets_sha256": paths["qwen_assets_sha256"], + "starvla_checkout": str(paths["source_dir"]), + "starvla_revision": paths["source_revision"], + }, + "runtime": { + **_runtime_record( + torch, transformers, str(next(framework.parameters()).device) + ), + "qwen-vl-utils": _distribution_version("qwen-vl-utils"), + }, + "determinism": { + "seed": 0, + "torch_deterministic_algorithms": True, + "cublas_workspace_config": os.environ.get("CUBLAS_WORKSPACE_CONFIG"), + "allow_tf32": False, + "attention_implementation": "sdpa", + }, + "input": { + "task": task, + "unnorm_key": unnorm_key, + "state": [], + "images": image_records, + "processed_images": [ + { + "index": index, + "mode": image.mode, + "size": list(image.size), + "pixel_sha256": _image_pixel_sha256(image), + } + for index, image in enumerate(captures["processed_images"]) + ], + }, + "normalization": legacy_normalization_contract( + paths["norm_stats"], unnorm_key + ), + "model_contract": { + "framework_class": f"{type(framework).__module__}.{type(framework).__name__}", + "action_token": ACTION_TOKEN, + "action_token_id": action_token_id, + "action_horizon": int(framework.chunk_len), + "action_dim": int(framework.action_model.action_dim), + "qwen_hidden_dim": int(framework.qwen_vl_interface.model.config.hidden_size), + "policy_input_dtype": captures["policy_input_dtype"], + "processor_class": type( + framework.qwen_vl_interface.processor + ).__name__, + "image_processor_class": type( + framework.qwen_vl_interface.processor.image_processor + ).__name__, + "image_patch_size": patch_size, + "image_spatial_merge_size": spatial_merge_size, + "image_grid_thw": grid_thw, + "image_resized_size": resized_size, + "merged_image_token_count": merged_image_token_count, + }, + "prompt": { + "framework_instruction": framework_instruction, + "model_instruction": model_instruction, + "rendered_chat_template": rendered_prompt, + }, + "tokens": { + "input_ids": input_ids.tolist(), + "token_strings": token_strings, + "all_action_token_positions": all_positions, + "selected_action_token_positions": selected_positions, + }, + "outputs": { + "normalized_actions": arrays["normalized_actions"].tolist(), + "unnormalized_actions": arrays["unnormalized_actions"].tolist(), + }, + "action_gate": { + "metric": "full_tensor_global_relative_l2", + "operator": "<=", + "limit": ACTION_RELATIVE_L2_LIMIT, + "required_outputs": ["normalized_actions", "unnormalized_actions"], + }, + "artifacts": { + "tensors": { + "path": tensors_path.name, + "size": tensors_path.stat().st_size, + "sha256": sha256_file(tensors_path), + "arrays": array_records, + } + }, + } + manifest_path = staging / "golden.json" + manifest_path.write_text( + json.dumps(manifest, ensure_ascii=False, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + Path(temporary).replace(output_dir) + return output_dir / "golden.json" + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description=( + "Generate an auditable local-Python golden from the official " + "StarVLA Qwen2.5-VL OFT .pt checkpoint." + ) + ) + parser.add_argument("--checkpoint", required=True, type=Path) + parser.add_argument( + "--qwen-model", + required=True, + type=Path, + help="Local Qwen2.5-VL-3B-Instruct topology and processor directory", + ) + parser.add_argument( + "--starvla-source", + type=Path, + default=Path("ckpts/starvla/source/starvla"), + ) + parser.add_argument( + "--expected-checkpoint-sha256", + default=OFFICIAL_CHECKPOINT_SHA256, + ) + parser.add_argument( + "--expected-checkpoint-size", + default=OFFICIAL_CHECKPOINT_SIZE, + type=int, + ) + parser.add_argument("--expected-source-revision") + parser.add_argument("--image", action="append", default=[], type=Path) + parser.add_argument("--task") + parser.add_argument("--unnorm-key", choices=tuple(LEGACY_UNNORM_PROFILES)) + parser.add_argument("--output-dir", type=Path) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--preflight-only", action="store_true") + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + args = build_parser().parse_args(argv) + _require_isolated_python() + if not args.preflight_only: + missing = [ + name + for name, value in ( + ("--image", args.image), + ("--task", args.task), + ("--unnorm-key", args.unnorm_key), + ("--output-dir", args.output_dir), + ) + if not value + ] + if missing: + raise StarVLAError("golden generation requires " + ", ".join(missing)) + if not args.task.strip() or not args.unnorm_key.strip(): + raise StarVLAError("--task and --unnorm-key must not be empty") + if len(args.image) != 1: + raise StarVLAError( + "Qwen2.5 OFT Bridge golden generation requires exactly one --image" + ) + + paths = validate_local_inputs( + checkpoint=args.checkpoint, + qwen_model=args.qwen_model, + source_dir=args.starvla_source, + expected_checkpoint_sha256=args.expected_checkpoint_sha256, + expected_checkpoint_size=args.expected_checkpoint_size, + expected_source_revision=args.expected_source_revision, + ) + try: + import torch + import transformers + except ImportError as exc: + raise StarVLAError(f"official StarVLA runtime dependency is missing: {exc}") from exc + validate_runtime_versions( + torch_version=torch.__version__, + torchvision_version=_distribution_version("torchvision"), + transformers_version=transformers.__version__, + numpy_version=np.__version__, + ) + qwen_vl_utils_version = _distribution_version("qwen-vl-utils") + if qwen_vl_utils_version != EXPECTED_QWEN_VL_UTILS_VERSION: + raise StarVLAError( + "official Qwen2.5 OFT oracle requires qwen-vl-utils " + f"{EXPECTED_QWEN_VL_UTILS_VERSION}, got {qwen_vl_utils_version}" + ) + _configure_determinism(torch, seed=0, device=args.device) + if args.preflight_only: + print( + "Qwen2.5 OFT local .pt preflight passed: " + f"{paths['checkpoint']} ({paths['checkpoint_sha256']})" + ) + return 0 + + images, image_records = _load_images(args.image) + framework, config = load_official_framework(paths, device=args.device) + captures = run_official_forward(framework, images=images, task=args.task) + + normalized = captures["normalized_actions"] + unnormalized = unnormalize_legacy_actions( + normalized, paths["norm_stats"], args.unnorm_key + ) + + manifest = write_golden( + output_dir=args.output_dir, + paths=paths, + framework=framework, + config=config, + image_paths=args.image, + source_image_records=image_records, + task=args.task, + unnorm_key=args.unnorm_key, + captures=captures, + unnormalized=np.ascontiguousarray(unnormalized, dtype=np.float32), + ) + print(f"Wrote StarVLA Qwen2.5 OFT local-Python golden: {manifest}") + return 0 + + +if __name__ == "__main__": + try: + raise SystemExit(main()) + except StarVLAError as exc: + raise SystemExit(f"error: {exc}") from exc From bb5093cab47cd319edc76a00e3676ce4d9e44333 Mon Sep 17 00:00:00 2001 From: JJJYmmm <1650675829@qq.com> Date: Mon, 10 Aug 2026 12:42:56 +0800 Subject: [PATCH 04/11] starvla: add Qwen3-VL PI v3 policy --- src/models/starvla/pi_v3_policy.cpp | 955 +++++++++ src/models/starvla/pi_v3_policy.h | 98 + tests/starvla/test_starvla_pi_v3_golden.py | 36 + .../starvla/generate_starvla_pi_v3_golden.py | 1831 +++++++++++++++++ .../starvla/pi_v3_golden_constraints.txt | 14 + 5 files changed, 2934 insertions(+) create mode 100644 src/models/starvla/pi_v3_policy.cpp create mode 100644 src/models/starvla/pi_v3_policy.h create mode 100644 tests/starvla/test_starvla_pi_v3_golden.py create mode 100644 tools/hf2gguf/starvla/generate_starvla_pi_v3_golden.py create mode 100644 tools/hf2gguf/starvla/pi_v3_golden_constraints.txt diff --git a/src/models/starvla/pi_v3_policy.cpp b/src/models/starvla/pi_v3_policy.cpp new file mode 100644 index 0000000..8cc1fd1 --- /dev/null +++ b/src/models/starvla/pi_v3_policy.cpp @@ -0,0 +1,955 @@ +#include "models/starvla/pi_v3_policy.h" + +#include "ggml-backend.h" +#include "ggml.h" +#include "gguf.h" +#include "models/ggml_backend.h" +#include "models/gguf_loader.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace robotcpp::starvla { + +namespace { + +constexpr size_t kGraphSize = 32768; +constexpr int kKQMaskPad = 32; +constexpr int kReleasedLayerCount = 36; + +struct PIV3BlockWeights { + ggml_tensor * ada_norm_weight = nullptr; + ggml_tensor * ada_norm_bias = nullptr; + ggml_tensor * query_weight = nullptr; + ggml_tensor * query_bias = nullptr; + ggml_tensor * key_weight = nullptr; + ggml_tensor * key_bias = nullptr; + ggml_tensor * value_weight = nullptr; + ggml_tensor * value_bias = nullptr; + ggml_tensor * attention_output_weight = nullptr; + ggml_tensor * attention_output_bias = nullptr; + ggml_tensor * feed_forward_input_weight = nullptr; + ggml_tensor * feed_forward_input_bias = nullptr; + ggml_tensor * feed_forward_output_weight = nullptr; + ggml_tensor * feed_forward_output_bias = nullptr; +}; + +struct PIV3ProjectorWeights { + ggml_tensor * norm_weight = nullptr; + ggml_tensor * norm_bias = nullptr; + ggml_tensor * projection_weight = nullptr; + ggml_tensor * projection_bias = nullptr; +}; + +struct PIV3Weights { + ggml_tensor * timestep_input_weight = nullptr; + ggml_tensor * timestep_input_bias = nullptr; + ggml_tensor * timestep_output_weight = nullptr; + ggml_tensor * timestep_output_bias = nullptr; + std::vector blocks; + std::vector projectors; + ggml_tensor * action_input_weight = nullptr; + ggml_tensor * action_input_bias = nullptr; + ggml_tensor * action_time_mix_weight = nullptr; + ggml_tensor * action_time_mix_bias = nullptr; + ggml_tensor * action_output_weight = nullptr; + ggml_tensor * action_output_bias = nullptr; + ggml_tensor * velocity_input_weight = nullptr; + ggml_tensor * velocity_input_bias = nullptr; + ggml_tensor * velocity_output_weight = nullptr; + ggml_tensor * velocity_output_bias = nullptr; + ggml_tensor * future_tokens = nullptr; + ggml_tensor * action_position = nullptr; +}; + +int require_key(gguf_context * gguf, const char * key, gguf_type type) { + const int index = gguf_find_key(gguf, key); + if (index < 0) { + throw std::runtime_error(std::string("missing required StarVLA GGUF metadata: ") + key); + } + if (gguf_get_kv_type(gguf, index) != type) { + throw std::runtime_error(std::string("invalid StarVLA GGUF metadata type: ") + key); + } + return index; +} + +std::string require_string(gguf_context * gguf, const char * key) { + return gguf_get_val_str(gguf, require_key(gguf, key, GGUF_TYPE_STRING)); +} + +int require_i32(gguf_context * gguf, const char * key) { + return gguf_get_val_i32(gguf, require_key(gguf, key, GGUF_TYPE_INT32)); +} + +float require_f32(gguf_context * gguf, const char * key) { + return gguf_get_val_f32(gguf, require_key(gguf, key, GGUF_TYPE_FLOAT32)); +} + +bool require_bool(gguf_context * gguf, const char * key) { + return gguf_get_val_bool(gguf, require_key(gguf, key, GGUF_TYPE_BOOL)); +} + +int require_array(gguf_context * gguf, const char * key, gguf_type element_type) { + const int index = require_key(gguf, key, GGUF_TYPE_ARRAY); + if (gguf_get_arr_type(gguf, index) != element_type) { + throw std::runtime_error(std::string("invalid StarVLA GGUF array element type: ") + key); + } + return index; +} + +std::vector require_string_array(gguf_context * gguf, const char * key) { + const int index = require_array(gguf, key, GGUF_TYPE_STRING); + const size_t count = gguf_get_arr_n(gguf, index); + std::vector result; + result.reserve(count); + for (size_t i = 0; i < count; ++i) { + result.emplace_back(gguf_get_arr_str(gguf, index, i)); + } + return result; +} + +std::vector require_i32_array(gguf_context * gguf, const char * key) { + const int index = require_array(gguf, key, GGUF_TYPE_INT32); + const size_t count = gguf_get_arr_n(gguf, index); + const auto * data = static_cast(gguf_get_arr_data(gguf, index)); + if (data == nullptr && count != 0) { + throw std::runtime_error(std::string("missing StarVLA GGUF array data: ") + key); + } + return count == 0 ? std::vector() : std::vector(data, data + count); +} + +std::vector require_f32_array(gguf_context * gguf, const char * key) { + const int index = require_array(gguf, key, GGUF_TYPE_FLOAT32); + const size_t count = gguf_get_arr_n(gguf, index); + const auto * data = static_cast(gguf_get_arr_data(gguf, index)); + if (data == nullptr && count != 0) { + throw std::runtime_error(std::string("missing StarVLA GGUF array data: ") + key); + } + return count == 0 ? std::vector() : std::vector(data, data + count); +} + +std::vector require_bool_array(gguf_context * gguf, const char * key) { + const int index = require_array(gguf, key, GGUF_TYPE_BOOL); + const size_t count = gguf_get_arr_n(gguf, index); + const auto * data = static_cast(gguf_get_arr_data(gguf, index)); + if (data == nullptr && count != 0) { + throw std::runtime_error(std::string("missing StarVLA GGUF array data: ") + key); + } + std::vector result(count); + for (size_t i = 0; i < count; ++i) { + result[i] = data[i] != 0 ? 1 : 0; + } + return result; +} + +std::string profile_key(int profile_index, const char * suffix) { + return "starvla.normalization.profile." + std::to_string(profile_index) + "." + suffix; +} + +bool has_shape(const ggml_tensor * tensor, std::initializer_list expected) { + if (tensor == nullptr || static_cast(ggml_n_dims(tensor)) != expected.size()) { + return false; + } + size_t dimension = 0; + for (const int64_t value : expected) { + if (tensor->ne[dimension++] != value) { + return false; + } + } + return true; +} + +const char * mode_name(backend_mode mode) { + switch (mode) { + case backend_mode::cpu: + return "cpu"; + case backend_mode::cuda: + return "cuda"; + case backend_mode::metal: + return "metal"; + } + return "unknown"; +} + +std::vector integer_range(int first, int count) { + std::vector result(static_cast(count)); + for (int index = 0; index < count; ++index) { + result[static_cast(index)] = first + index; + } + return result; +} + +class PIV3GGUFLoader final : public gguf_loader { + public: + PIV3GGUFLoader(PIV3PolicyConfig & config, PIV3Weights & weights) : config_(config), weights_(weights) {} + + protected: + bool parse_metadata(gguf_context * gguf) override { + if (require_string(gguf, "general.architecture") != "starvla-policy" || + require_i32(gguf, "starvla.schema_version") != 1 || + require_string(gguf, "starvla.framework") != "pi_v3") { + throw std::runtime_error("GGUF is not a supported StarVLA PI-v3 policy"); + } + + config_.backbone_arch = require_string(gguf, "starvla.backbone.arch"); + if (config_.backbone_arch != "qwen3_vl") { + throw std::runtime_error("StarVLA PI-v3 requires a Qwen3-VL backbone"); + } + config_.bundle_uuid = require_string(gguf, "starvla.bundle.uuid"); + config_.text_filename = require_string(gguf, "starvla.component.text.filename"); + config_.mmproj_filename = require_string(gguf, "starvla.component.mmproj.filename"); + if (config_.bundle_uuid.empty() || config_.text_filename.empty() || + config_.mmproj_filename.empty()) { + throw std::runtime_error("StarVLA PI-v3 bundle metadata is incomplete"); + } + + config_.qwen_hidden_dim = require_i32(gguf, "starvla.qwen.hidden_size"); + const int embedding_key = gguf_find_key(gguf, "starvla.qwen.input_embedding_size"); + config_.qwen_input_embedding_dim = + embedding_key < 0 ? 4 * config_.qwen_hidden_dim + : require_i32(gguf, "starvla.qwen.input_embedding_size"); + config_.qwen_layer_count = require_i32(gguf, "starvla.qwen.layer_count"); + config_.qwen_vocab_size = require_i32(gguf, "starvla.qwen.vocab_size"); + config_.cot_template = require_string(gguf, "starvla.prompt.cot_template"); + + config_.image_count = require_i32(gguf, "starvla.image.count"); + config_.image_names = require_string_array(gguf, "starvla.image.names"); + config_.image_processor_min_pixels = + require_i32(gguf, "starvla.image.processor_min_pixels"); + config_.image_processor_max_pixels = + require_i32(gguf, "starvla.image.processor_max_pixels"); + config_.image_patch_size = require_i32(gguf, "starvla.image.patch_size"); + config_.image_spatial_merge_size = + require_i32(gguf, "starvla.image.spatial_merge_size"); + config_.image_min_token_count = + require_i32(gguf, "starvla.image.min_token_count"); + config_.image_max_token_count = + require_i32(gguf, "starvla.image.max_token_count"); + + config_.dit_width = require_i32(gguf, "starvla.pi_v3.dit_width"); + config_.block_count = require_i32(gguf, "starvla.pi_v3.block_count"); + config_.projector_count = require_i32(gguf, "starvla.pi_v3.projector_count"); + config_.attention_head_count = + require_i32(gguf, "starvla.pi_v3.attention_head_count"); + config_.attention_head_dim = + require_i32(gguf, "starvla.pi_v3.attention_head_dim"); + config_.feed_forward_dim = + require_i32(gguf, "starvla.pi_v3.feed_forward_dim"); + config_.mlp_hidden_dim = + require_i32(gguf, "starvla.pi_v3.mlp_hidden_dimension"); + config_.future_token_count = + require_i32(gguf, "starvla.pi_v3.future_token_count"); + config_.action_position_count = + require_i32(gguf, "starvla.pi_v3.action_position_count"); + config_.no_state_sequence_length = + require_i32(gguf, "starvla.pi_v3.no_state_sequence_length"); + config_.timestep_projection_dim = + require_i32(gguf, "starvla.pi_v3.timestep_projection_dim"); + config_.num_timestep_buckets = + require_i32(gguf, "starvla.pi_v3.num_timestep_buckets"); + config_.num_inference_timesteps = + require_i32(gguf, "starvla.pi_v3.num_inference_timesteps"); + config_.ada_norm_epsilon = require_f32(gguf, "starvla.pi_v3.ada_norm_epsilon"); + config_.projector_norm_epsilon = + require_f32(gguf, "starvla.pi_v3.projector_norm_epsilon"); + config_.euler_dt = require_f32(gguf, "starvla.pi_v3.euler_dt"); + config_.action_dim = require_i32(gguf, "starvla.action.dimension"); + config_.horizon = require_i32(gguf, "starvla.action.horizon"); + + const bool valid = + config_.qwen_hidden_dim > 0 && config_.qwen_input_embedding_dim > 0 && + config_.qwen_layer_count == kReleasedLayerCount && config_.qwen_vocab_size > 0 && + !config_.cot_template.empty() && config_.image_count > 0 && + config_.image_names.size() == static_cast(config_.image_count) && + config_.image_processor_min_pixels > 0 && + config_.image_processor_max_pixels >= config_.image_processor_min_pixels && + config_.image_patch_size > 0 && config_.image_spatial_merge_size > 0 && + config_.image_min_token_count > 0 && + config_.image_max_token_count >= config_.image_min_token_count && + config_.dit_width > 0 && config_.block_count == kReleasedLayerCount && + config_.projector_count == config_.block_count && + config_.attention_head_count > 0 && config_.attention_head_dim > 0 && + config_.attention_head_count * config_.attention_head_dim == config_.dit_width && + config_.feed_forward_dim > 0 && config_.mlp_hidden_dim > 0 && + config_.action_dim > 0 && config_.horizon > 0 && + config_.future_token_count > 0 && + config_.action_position_count >= config_.horizon && + config_.no_state_sequence_length == + config_.future_token_count + config_.horizon && + config_.timestep_projection_dim >= 4 && + config_.timestep_projection_dim % 2 == 0 && + config_.num_timestep_buckets > 0 && config_.num_inference_timesteps == 4 && + config_.ada_norm_epsilon > 0.0f && config_.projector_norm_epsilon > 0.0f && + config_.euler_dt > 0.0f; + if (!valid) { + throw std::runtime_error("StarVLA PI-v3 metadata has incompatible dimensions"); + } + + config_.qwen_hidden_tuple_indices = integer_range(1, config_.qwen_layer_count); + config_.timestep_ids.resize(static_cast(config_.num_inference_timesteps)); + for (int step = 0; step < config_.num_inference_timesteps; ++step) { + config_.timestep_ids[static_cast(step)] = + step * config_.num_timestep_buckets / config_.num_inference_timesteps; + } + + NormalizationConfig & normalization = config_.normalization; + normalization.clip_actions = require_bool(gguf, "starvla.normalization.clip_actions"); + normalization.binary_threshold = + require_f32(gguf, "starvla.normalization.binary_threshold"); + normalization.binary_comparison = + require_string(gguf, "starvla.normalization.binary_comparison"); + normalization.continuous_dimensions = + require_i32_array(gguf, "starvla.action.continuous_dimensions"); + normalization.binary_dimensions = + require_i32_array(gguf, "starvla.action.binary_dimensions"); + const int profile_count = + require_i32(gguf, "starvla.normalization.profile_count"); + const std::vector keys = + require_string_array(gguf, "starvla.normalization.profile_keys"); + if (profile_count <= 0 || keys.size() != static_cast(profile_count)) { + throw std::runtime_error("StarVLA PI-v3 normalization profiles are inconsistent"); + } + normalization.profiles.clear(); + normalization.profiles.reserve(static_cast(profile_count)); + for (int index = 0; index < profile_count; ++index) { + NormalizationProfile profile; + profile.key = require_string(gguf, profile_key(index, "key").c_str()); + profile.action_q01 = + require_f32_array(gguf, profile_key(index, "action_q01").c_str()); + profile.action_q99 = + require_f32_array(gguf, profile_key(index, "action_q99").c_str()); + profile.action_mask = + require_bool_array(gguf, profile_key(index, "action_mask").c_str()); + if (profile.key != keys[static_cast(index)]) { + throw std::runtime_error("StarVLA PI-v3 normalization profile order is inconsistent"); + } + normalization.profiles.push_back(std::move(profile)); + } + std::string normalization_error; + if (!validate_normalization_config(normalization, config_.action_dim, + normalization_error)) { + throw std::runtime_error(normalization_error); + } + return true; + } + + bool bind_tensors(ggml_context * ctx_data) override { + auto bind = [&](ggml_tensor *& destination, const std::string & name) { + destination = require_tensor(ctx_data, name); + }; + bind(weights_.timestep_input_weight, "starvla.policy.pi_v3.timestep.input.weight"); + bind(weights_.timestep_input_bias, "starvla.policy.pi_v3.timestep.input.bias"); + bind(weights_.timestep_output_weight, "starvla.policy.pi_v3.timestep.output.weight"); + bind(weights_.timestep_output_bias, "starvla.policy.pi_v3.timestep.output.bias"); + weights_.blocks.clear(); + weights_.blocks.reserve(static_cast(config_.block_count)); + for (int block = 0; block < config_.block_count; ++block) { + const std::string prefix = "starvla.policy.pi_v3.block." + std::to_string(block) + "."; + PIV3BlockWeights current; + bind(current.ada_norm_weight, prefix + "ada_norm.weight"); + bind(current.ada_norm_bias, prefix + "ada_norm.bias"); + bind(current.query_weight, prefix + "attention.query.weight"); + bind(current.query_bias, prefix + "attention.query.bias"); + bind(current.key_weight, prefix + "attention.key.weight"); + bind(current.key_bias, prefix + "attention.key.bias"); + bind(current.value_weight, prefix + "attention.value.weight"); + bind(current.value_bias, prefix + "attention.value.bias"); + bind(current.attention_output_weight, prefix + "attention.output.weight"); + bind(current.attention_output_bias, prefix + "attention.output.bias"); + bind(current.feed_forward_input_weight, prefix + "feed_forward.input.weight"); + bind(current.feed_forward_input_bias, prefix + "feed_forward.input.bias"); + bind(current.feed_forward_output_weight, prefix + "feed_forward.output.weight"); + bind(current.feed_forward_output_bias, prefix + "feed_forward.output.bias"); + weights_.blocks.push_back(current); + } + weights_.projectors.clear(); + weights_.projectors.reserve(static_cast(config_.projector_count)); + for (int projector = 0; projector < config_.projector_count; ++projector) { + const std::string prefix = "starvla.policy.pi_v3.projector." + + std::to_string(projector) + "."; + PIV3ProjectorWeights current; + bind(current.norm_weight, prefix + "norm.weight"); + bind(current.norm_bias, prefix + "norm.bias"); + bind(current.projection_weight, prefix + "projection.weight"); + bind(current.projection_bias, prefix + "projection.bias"); + weights_.projectors.push_back(current); + } + bind(weights_.action_input_weight, "starvla.policy.pi_v3.action.input.weight"); + bind(weights_.action_input_bias, "starvla.policy.pi_v3.action.input.bias"); + bind(weights_.action_time_mix_weight, "starvla.policy.pi_v3.action.time_mix.weight"); + bind(weights_.action_time_mix_bias, "starvla.policy.pi_v3.action.time_mix.bias"); + bind(weights_.action_output_weight, "starvla.policy.pi_v3.action.output.weight"); + bind(weights_.action_output_bias, "starvla.policy.pi_v3.action.output.bias"); + bind(weights_.velocity_input_weight, "starvla.policy.pi_v3.velocity.input.weight"); + bind(weights_.velocity_input_bias, "starvla.policy.pi_v3.velocity.input.bias"); + bind(weights_.velocity_output_weight, "starvla.policy.pi_v3.velocity.output.weight"); + bind(weights_.velocity_output_bias, "starvla.policy.pi_v3.velocity.output.bias"); + bind(weights_.future_tokens, "starvla.policy.pi_v3.future_tokens.weight"); + bind(weights_.action_position, "starvla.policy.pi_v3.action_position.weight"); + + const int width = config_.dit_width; + if (!has_shape(weights_.timestep_input_weight, {config_.timestep_projection_dim, width}) || + !has_shape(weights_.timestep_input_bias, {width}) || + !has_shape(weights_.timestep_output_weight, {width, width}) || + !has_shape(weights_.timestep_output_bias, {width}) || + !has_shape(weights_.action_input_weight, {config_.action_dim, width}) || + !has_shape(weights_.action_input_bias, {width}) || + !has_shape(weights_.action_time_mix_weight, {2 * width, width}) || + !has_shape(weights_.action_time_mix_bias, {width}) || + !has_shape(weights_.action_output_weight, {width, width}) || + !has_shape(weights_.action_output_bias, {width}) || + !has_shape(weights_.velocity_input_weight, {width, config_.mlp_hidden_dim}) || + !has_shape(weights_.velocity_input_bias, {config_.mlp_hidden_dim}) || + !has_shape(weights_.velocity_output_weight, {config_.mlp_hidden_dim, config_.action_dim}) || + !has_shape(weights_.velocity_output_bias, {config_.action_dim}) || + !has_shape(weights_.future_tokens, {width, config_.future_token_count}) || + !has_shape(weights_.action_position, {width, config_.action_position_count})) { + throw std::runtime_error("StarVLA PI_v3 non-block tensor has an incompatible ggml shape"); + } + for (const PIV3BlockWeights & current : weights_.blocks) { + if (!has_shape(current.ada_norm_weight, {width, 2 * width}) || + !has_shape(current.ada_norm_bias, {2 * width}) || + !has_shape(current.query_weight, {width, width}) || + !has_shape(current.query_bias, {width}) || + !has_shape(current.key_weight, {width, width}) || + !has_shape(current.key_bias, {width}) || + !has_shape(current.value_weight, {width, width}) || + !has_shape(current.value_bias, {width}) || + !has_shape(current.attention_output_weight, {width, width}) || + !has_shape(current.attention_output_bias, {width}) || + !has_shape(current.feed_forward_input_weight, {width, config_.feed_forward_dim}) || + !has_shape(current.feed_forward_input_bias, {config_.feed_forward_dim}) || + !has_shape(current.feed_forward_output_weight, {config_.feed_forward_dim, width}) || + !has_shape(current.feed_forward_output_bias, {width})) { + throw std::runtime_error("StarVLA PI_v3 transformer block tensor has an incompatible ggml shape"); + } + } + for (const PIV3ProjectorWeights & current : weights_.projectors) { + if (!has_shape(current.norm_weight, {config_.qwen_hidden_dim}) || + !has_shape(current.norm_bias, {config_.qwen_hidden_dim}) || + !has_shape(current.projection_weight, {config_.qwen_hidden_dim, width}) || + !has_shape(current.projection_bias, {width})) { + throw std::runtime_error("StarVLA PI_v3 projector tensor has an incompatible ggml shape"); + } + } + return true; + } + + private: + PIV3PolicyConfig & config_; + PIV3Weights & weights_; +}; + +std::vector timestep_projection_table(const PIV3PolicyConfig & config) { + const int dim = config.timestep_projection_dim; + const int half = dim / 2; + const float denominator = static_cast(half - 1); + std::vector table(static_cast(dim) * 4, 0.0f); + for (int step = 0; step < 4; ++step) { + const float timestep = static_cast(config.timestep_ids[static_cast(step)]); + float * row = table.data() + static_cast(step) * dim; + for (int index = 0; index < half; ++index) { + const float frequency = std::exp(-std::log(10000.0f) * static_cast(index) / denominator); + const float angle = timestep * frequency; + row[index] = std::cos(angle); + row[index + half] = std::sin(angle); + } + } + return table; +} + +std::vector action_time_table(const PIV3PolicyConfig & config) { + const int dim = config.dit_width; + const int half = dim / 2; + const float denominator = static_cast(half); + std::vector table(static_cast(dim) * 4, 0.0f); + for (int step = 0; step < 4; ++step) { + const float timestep = static_cast(config.timestep_ids[static_cast(step)]); + float * row = table.data() + static_cast(step) * dim; + for (int index = 0; index < half; ++index) { + const float frequency = std::exp(-std::log(10000.0f) * static_cast(index) / denominator); + const float angle = timestep * frequency; + row[index] = std::sin(angle); + row[index + half] = std::cos(angle); + } + } + return table; +} + +} // namespace + +struct PIV3Policy::Impl { + PIV3PolicyConfig config; + PIV3Weights weights; + gguf_load_result loaded; + ggml_backend_t backend_cpu = nullptr; + std::vector backends; + ggml_backend_sched_t scheduler = nullptr; + backend_buft_policy buft_policy; + backend_mode mode = backend_mode::cpu; + int n_threads = 0; + int verbosity = 0; + ggml_context * graph_context = nullptr; + ggml_cgraph * graph = nullptr; + ggml_tensor * hidden_input = nullptr; + ggml_tensor * cross_mask_input = nullptr; + ggml_tensor * noise_input = nullptr; + ggml_tensor * timestep_projection_input = nullptr; + ggml_tensor * action_time_input = nullptr; + ggml_tensor * scalar_one_input = nullptr; + ggml_tensor * output = nullptr; + size_t conditioning_token_count = 0; + std::vector timestep_table; + std::vector action_table; + ~Impl() { + clear_graph(); + if (scheduler != nullptr) { + ggml_backend_sched_synchronize(scheduler); + ggml_backend_sched_free(scheduler); + scheduler = nullptr; + } + if (loaded.model_buffer != nullptr) { + ggml_backend_buffer_free(loaded.model_buffer); + loaded.model_buffer = nullptr; + } + if (loaded.ctx_data != nullptr) { + ggml_free(loaded.ctx_data); + loaded.ctx_data = nullptr; + } + if (loaded.gguf != nullptr) { + gguf_free(loaded.gguf); + loaded.gguf = nullptr; + } + for (ggml_backend_t backend : backends) { + if (backend != nullptr) { + ggml_backend_free(backend); + } + } + backends.clear(); + backend_cpu = nullptr; + } + + void clear_graph() { + if (scheduler != nullptr) { + ggml_backend_sched_synchronize(scheduler); + ggml_backend_sched_reset(scheduler); + } + if (graph_context != nullptr) { + ggml_free(graph_context); + graph_context = nullptr; + } + graph = nullptr; + hidden_input = nullptr; + cross_mask_input = nullptr; + noise_input = nullptr; + timestep_projection_input = nullptr; + action_time_input = nullptr; + scalar_one_input = nullptr; + output = nullptr; + conditioning_token_count = 0; + } + + void build_graph(size_t token_count) { + clear_graph(); + if (token_count == 0 || token_count > static_cast(std::numeric_limits::max())) { + throw std::runtime_error("invalid StarVLA PI_v3 conditioning token count"); + } + + ggml_init_params params{}; + params.mem_size = kGraphSize * ggml_tensor_overhead() + + ggml_graph_overhead_custom(kGraphSize, false); + params.mem_buffer = nullptr; + params.no_alloc = true; + graph_context = ggml_init(params); + if (graph_context == nullptr) { + throw std::runtime_error("failed to initialize StarVLA PI_v3 graph context"); + } + + const int width = config.dit_width; + const int heads = config.attention_head_count; + const int head_dim = config.attention_head_dim; + const int sequence_length = config.no_state_sequence_length; + const int mask_queries = GGML_PAD(sequence_length, kKQMaskPad); + + hidden_input = ggml_new_tensor_3d(graph_context, GGML_TYPE_F32, config.qwen_hidden_dim, + static_cast(token_count), config.qwen_layer_count); + cross_mask_input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, + static_cast(token_count), mask_queries); + noise_input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, + config.action_dim, config.horizon); + timestep_projection_input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, + config.timestep_projection_dim, 4); + action_time_input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, width, 4); + scalar_one_input = ggml_new_tensor_1d(graph_context, GGML_TYPE_F32, 1); + if (hidden_input == nullptr || cross_mask_input == nullptr || noise_input == nullptr || + timestep_projection_input == nullptr || action_time_input == nullptr || + scalar_one_input == nullptr) { + throw std::runtime_error("failed to create StarVLA PI_v3 graph inputs"); + } + ggml_set_name(hidden_input, "starvla_pi_v3_qwen_hidden_states"); + ggml_set_name(cross_mask_input, "starvla_pi_v3_qwen_attention_mask"); + ggml_set_name(noise_input, "starvla_pi_v3_initial_noise"); + ggml_set_name(timestep_projection_input, "starvla_pi_v3_timestep_projection_table"); + ggml_set_name(action_time_input, "starvla_pi_v3_action_time_table"); + ggml_set_name(scalar_one_input, "starvla_pi_v3_scalar_one"); + ggml_set_input(hidden_input); + ggml_set_input(cross_mask_input); + ggml_set_input(noise_input); + ggml_set_input(timestep_projection_input); + ggml_set_input(action_time_input); + ggml_set_input(scalar_one_input); + + auto f32 = [&](ggml_tensor * tensor) { + return tensor->type == GGML_TYPE_F32 ? tensor : + ggml_cast(graph_context, tensor, GGML_TYPE_F32); + }; + auto bf16_roundtrip = [&](ggml_tensor * tensor) { + return ggml_cast(graph_context, + ggml_cast(graph_context, tensor, GGML_TYPE_BF16), + GGML_TYPE_F32); + }; + auto linear = [&](ggml_tensor * value, ggml_tensor * weight, ggml_tensor * bias) { + ggml_tensor * projected = ggml_mul_mat(graph_context, weight, value); + ggml_mul_mat_set_prec(projected, GGML_PREC_F32); + return ggml_add(graph_context, projected, f32(bias)); + }; + auto projector_linear = [&](ggml_tensor * value, ggml_tensor * weight, + ggml_tensor * bias) { + ggml_tensor * bf16_value = ggml_cast(graph_context, value, GGML_TYPE_BF16); + ggml_tensor * bf16_weight = ggml_cast(graph_context, weight, GGML_TYPE_BF16); + ggml_tensor * projected = + ggml_mul_mat(graph_context, bf16_weight, bf16_value); + ggml_mul_mat_set_prec(projected, GGML_PREC_F32); + projected = ggml_add(graph_context, projected, bf16_roundtrip(bias)); + return bf16_roundtrip(projected); + }; + auto ada_norm = [&](ggml_tensor * value, ggml_tensor * temb, + const PIV3BlockWeights & block) { + ggml_tensor * modulation = linear(ggml_silu(graph_context, temb), + block.ada_norm_weight, block.ada_norm_bias); + ggml_tensor * scale = ggml_view_1d(graph_context, modulation, width, 0); + ggml_tensor * shift = ggml_view_1d(graph_context, modulation, width, + static_cast(width) * sizeof(float)); + ggml_tensor * normalized = ggml_norm(graph_context, value, config.ada_norm_epsilon); + return ggml_add(graph_context, + ggml_mul(graph_context, normalized, + ggml_add(graph_context, scale, scalar_one_input)), + shift); + }; + auto attention = [&](ggml_tensor * query_source, ggml_tensor * key_value_source, + const PIV3BlockWeights & block) { + const int64_t query_count = query_source->ne[1]; + const int64_t key_value_count = key_value_source->ne[1]; + ggml_tensor * query = linear(query_source, block.query_weight, block.query_bias); + ggml_tensor * key = linear(key_value_source, block.key_weight, block.key_bias); + ggml_tensor * value = linear(key_value_source, block.value_weight, block.value_bias); + query = ggml_reshape_3d(graph_context, query, head_dim, heads, query_count); + key = ggml_reshape_3d(graph_context, key, head_dim, heads, key_value_count); + value = ggml_reshape_3d(graph_context, value, head_dim, heads, key_value_count); + query = ggml_permute(graph_context, query, 0, 2, 1, 3); + key = ggml_permute(graph_context, key, 0, 2, 1, 3); + value = ggml_cont(graph_context, ggml_permute(graph_context, value, 1, 2, 0, 3)); + ggml_tensor * scores = ggml_mul_mat(graph_context, key, query); + ggml_mul_mat_set_prec(scores, GGML_PREC_F32); + scores = ggml_soft_max_ext(graph_context, scores, cross_mask_input, + 1.0f / std::sqrt(static_cast(head_dim)), 0.0f); + ggml_tensor * attended = ggml_mul_mat(graph_context, value, scores); + ggml_mul_mat_set_prec(attended, GGML_PREC_F32); + attended = ggml_permute(graph_context, attended, 0, 2, 1, 3); + attended = ggml_cont_2d(graph_context, attended, width, query_count); + return linear(attended, block.attention_output_weight, block.attention_output_bias); + }; + + std::vector projected_hidden_states; + projected_hidden_states.reserve(static_cast(config.projector_count)); + for (int layer = 0; layer < config.projector_count; ++layer) { + const PIV3ProjectorWeights & projector = weights.projectors[static_cast(layer)]; + ggml_tensor * layer_hidden = ggml_view_2d( + graph_context, hidden_input, config.qwen_hidden_dim, + static_cast(token_count), hidden_input->nb[1], + static_cast(layer) * hidden_input->nb[2]); + layer_hidden = bf16_roundtrip(layer_hidden); + layer_hidden = ggml_norm(graph_context, layer_hidden, config.projector_norm_epsilon); + layer_hidden = ggml_mul(graph_context, layer_hidden, f32(projector.norm_weight)); + layer_hidden = ggml_add(graph_context, layer_hidden, f32(projector.norm_bias)); + ggml_tensor * projected = projector_linear( + layer_hidden, projector.projection_weight, + projector.projection_bias); + projected_hidden_states.push_back(projected); + } + + ggml_tensor * future = f32(weights.future_tokens); + ggml_tensor * position_view = ggml_view_2d( + graph_context, weights.action_position, width, config.horizon, + weights.action_position->nb[1], 0); + ggml_tensor * position = f32(position_view); + // Qwen/projector inference and torch.randn run at BF16 in the released + // script. The action head then enters CUDA autocast(float32). + ggml_tensor * actions = bf16_roundtrip(noise_input); + + for (int step = 0; step < 4; ++step) { + ggml_tensor * timestep_projection = ggml_view_1d( + graph_context, timestep_projection_input, config.timestep_projection_dim, + static_cast(step) * config.timestep_projection_dim * sizeof(float)); + ggml_tensor * temb = linear(timestep_projection, weights.timestep_input_weight, + weights.timestep_input_bias); + temb = ggml_silu(graph_context, temb); + temb = linear(temb, weights.timestep_output_weight, weights.timestep_output_bias); + + ggml_tensor * action_features = linear(actions, weights.action_input_weight, + weights.action_input_bias); + ggml_tensor * action_time = ggml_view_1d( + graph_context, action_time_input, width, + static_cast(step) * width * sizeof(float)); + action_time = ggml_repeat(graph_context, action_time, action_features); + action_features = ggml_concat(graph_context, action_features, action_time, 0); + action_features = linear(action_features, weights.action_time_mix_weight, + weights.action_time_mix_bias); + action_features = ggml_silu(graph_context, action_features); + action_features = linear(action_features, weights.action_output_weight, + weights.action_output_bias); + action_features = ggml_add(graph_context, action_features, position); + + ggml_tensor * hidden = ggml_concat(graph_context, future, action_features, 1); + for (int block_index = 0; block_index < config.block_count; ++block_index) { + const PIV3BlockWeights & block = weights.blocks[static_cast(block_index)]; + ggml_tensor * normalized = ada_norm(hidden, temb, block); + ggml_tensor * attended = attention( + normalized, projected_hidden_states[static_cast(block_index)], block); + hidden = ggml_add(graph_context, hidden, attended); + ggml_tensor * ff = ggml_norm(graph_context, hidden, config.ada_norm_epsilon); + ff = linear(ff, block.feed_forward_input_weight, block.feed_forward_input_bias); + ff = ggml_gelu(graph_context, ff); + ff = linear(ff, block.feed_forward_output_weight, block.feed_forward_output_bias); + hidden = ggml_add(graph_context, hidden, ff); + } + + // The released legacy sampler calls DiT with return_pre_output=true. + // norm_out/proj_out_1/proj_out_2 are therefore intentionally inactive. + hidden = ggml_relu(graph_context, + linear(hidden, weights.velocity_input_weight, + weights.velocity_input_bias)); + hidden = linear(hidden, weights.velocity_output_weight, + weights.velocity_output_bias); + ggml_tensor * velocity = ggml_view_2d( + graph_context, hidden, config.action_dim, config.horizon, hidden->nb[1], + static_cast(config.future_token_count) * hidden->nb[1]); + actions = ggml_add(graph_context, actions, + ggml_scale(graph_context, velocity, config.euler_dt)); + } + + output = actions; + ggml_set_name(output, "starvla_pi_v3_normalized_actions"); + ggml_set_output(output); + graph = ggml_new_graph_custom(graph_context, kGraphSize, false); + if (graph == nullptr) { + throw std::runtime_error("failed to create StarVLA PI_v3 graph"); + } + ggml_build_forward_expand(graph, output); + ggml_backend_sched_reset(scheduler); + if (!ggml_backend_sched_alloc_graph(scheduler, graph)) { + throw std::runtime_error("failed to allocate StarVLA PI_v3 graph"); + } + + conditioning_token_count = token_count; + } +}; + +PIV3Policy::PIV3Policy(std::unique_ptr impl) : impl_(std::move(impl)) {} + +PIV3Policy::~PIV3Policy() = default; + +std::unique_ptr PIV3Policy::load(const std::string & path, int n_threads, int verbosity, + std::string & error) { + error.clear(); + if (path.empty()) { + error = "StarVLA PI_v3 policy path is required"; + return nullptr; + } + + std::unique_ptr impl(new Impl()); + impl->n_threads = n_threads; + impl->verbosity = verbosity; + try { + backend_scheduler_config scheduler_config; + scheduler_config.max_nodes = static_cast(kGraphSize); + scheduler_config.parallel = false; + scheduler_config.op_offload = true; + backend_loader backend; + if (!backend.load(impl->backend_cpu, impl->backends, impl->scheduler, + impl->buft_policy, true, scheduler_config, verbosity)) { + error = "failed to initialize StarVLA PI_v3 backend: " + backend.error(); + return nullptr; + } + impl->mode = backend.mode(); + + PIV3GGUFLoader loader(impl->config, impl->weights); + if (!loader.load(path.c_str(), impl->buft_policy.model_buft, impl->loaded, verbosity)) { + error = loader.error(); + return nullptr; + } + if (impl->loaded.ctx_data == nullptr || impl->loaded.model_buffer == nullptr) { + error = "StarVLA PI_v3 policy GGUF has no tensors"; + return nullptr; + } + ggml_backend_buffer_set_usage(impl->loaded.model_buffer, + GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + impl->timestep_table = timestep_projection_table(impl->config); + impl->action_table = action_time_table(impl->config); + if (verbosity >= 1) { + std::fprintf(stderr, + "%s: backend=%s qwen=%d width=%d layers=%d horizon=%d action_dim=%d profiles=%zu\n", + __func__, mode_name(impl->mode), impl->config.qwen_hidden_dim, + impl->config.dit_width, impl->config.block_count, impl->config.horizon, + impl->config.action_dim, impl->config.normalization.profiles.size()); + } + } catch (const std::exception & exception) { + error = exception.what(); + return nullptr; + } + return std::unique_ptr(new PIV3Policy(std::move(impl))); +} + +bool PIV3Policy::evaluate(const float * qwen_hidden_states, size_t hidden_element_count, + const uint8_t * qwen_attention_mask, size_t mask_element_count, + const float * initial_noise, size_t noise_element_count, + std::vector & normalized_actions, std::string & error) { + return evaluate_internal(qwen_hidden_states, hidden_element_count, + qwen_attention_mask, mask_element_count, + initial_noise, noise_element_count, + normalized_actions, error); +} + +bool PIV3Policy::evaluate_internal( + const float * qwen_hidden_states, size_t hidden_element_count, + const uint8_t * qwen_attention_mask, size_t mask_element_count, + const float * initial_noise, size_t noise_element_count, + std::vector & normalized_actions, + std::string & error) { + normalized_actions.clear(); + error.clear(); + if (impl_ == nullptr || impl_->scheduler == nullptr) { + error = "StarVLA PI_v3 policy is not initialized"; + return false; + } + const size_t hidden_width = static_cast(impl_->config.qwen_hidden_dim); + const size_t layer_count = static_cast(impl_->config.qwen_layer_count); + if (qwen_hidden_states == nullptr || qwen_attention_mask == nullptr || initial_noise == nullptr || + mask_element_count == 0 || + mask_element_count > static_cast(std::numeric_limits::max()) || + mask_element_count > std::numeric_limits::max() / hidden_width || + mask_element_count * hidden_width > std::numeric_limits::max() / layer_count || + hidden_element_count != mask_element_count * hidden_width * layer_count) { + error = "StarVLA PI_v3 layerwise Qwen conditioning tensor or attention mask has an incompatible shape"; + return false; + } + const size_t expected_noise = + static_cast(impl_->config.horizon) * impl_->config.action_dim; + if (noise_element_count != expected_noise) { + error = "StarVLA PI_v3 initial-noise tensor has an incompatible shape"; + return false; + } + if (std::any_of(qwen_hidden_states, qwen_hidden_states + hidden_element_count, + [](float value) { return !std::isfinite(value); }) || + std::any_of(initial_noise, initial_noise + noise_element_count, + [](float value) { return !std::isfinite(value); })) { + error = "StarVLA PI_v3 conditioning and initial noise must be finite"; + return false; + } + bool has_valid_token = false; + for (size_t token = 0; token < mask_element_count; ++token) { + if (qwen_attention_mask[token] > 1) { + error = "StarVLA PI_v3 attention mask values must be zero or one"; + return false; + } + has_valid_token = has_valid_token || qwen_attention_mask[token] != 0; + } + if (!has_valid_token) { + error = "StarVLA PI_v3 attention mask must contain at least one valid token"; + return false; + } + + try { + if (impl_->graph == nullptr || + impl_->conditioning_token_count != mask_element_count) { + impl_->build_graph(mask_element_count); + } + } catch (const std::exception & exception) { + error = exception.what(); + return false; + } + + const int query_count = impl_->config.no_state_sequence_length; + const int padded_queries = GGML_PAD(query_count, kKQMaskPad); + std::vector additive_mask(mask_element_count * static_cast(padded_queries), + -std::numeric_limits::infinity()); + for (int query = 0; query < query_count; ++query) { + float * row = additive_mask.data() + static_cast(query) * mask_element_count; + for (size_t token = 0; token < mask_element_count; ++token) { + row[token] = qwen_attention_mask[token] != 0 + ? 0.0f + : -std::numeric_limits::infinity(); + } + } + + ggml_backend_tensor_set(impl_->hidden_input, qwen_hidden_states, 0, + hidden_element_count * sizeof(float)); + ggml_backend_tensor_set(impl_->cross_mask_input, additive_mask.data(), 0, + additive_mask.size() * sizeof(float)); + ggml_backend_tensor_set(impl_->noise_input, initial_noise, 0, + noise_element_count * sizeof(float)); + ggml_backend_tensor_set(impl_->timestep_projection_input, impl_->timestep_table.data(), 0, + impl_->timestep_table.size() * sizeof(float)); + ggml_backend_tensor_set(impl_->action_time_input, impl_->action_table.data(), 0, + impl_->action_table.size() * sizeof(float)); + const float one = 1.0f; + ggml_backend_tensor_set(impl_->scalar_one_input, &one, 0, sizeof(one)); + set_backend_threads(impl_->backends, impl_->n_threads); + if (ggml_backend_sched_graph_compute(impl_->scheduler, impl_->graph) != GGML_STATUS_SUCCESS) { + error = "StarVLA PI_v3 graph compute failed"; + return false; + } + + normalized_actions.resize(expected_noise); + ggml_backend_tensor_get(impl_->output, normalized_actions.data(), 0, + expected_noise * sizeof(float)); + if (std::any_of(normalized_actions.begin(), normalized_actions.end(), + [](float value) { return !std::isfinite(value); })) { + normalized_actions.clear(); + error = "StarVLA PI_v3 graph produced non-finite actions"; + return false; + } + return true; +} + +bool PIV3Policy::unnormalize(const std::vector & normalized_actions, + const std::string & profile_key_value, + std::vector & actions, std::string & error) const { + if (impl_ == nullptr) { + actions.clear(); + error = "StarVLA PI_v3 policy is not initialized"; + return false; + } + return denormalize_actions(impl_->config.normalization, profile_key_value, + normalized_actions, impl_->config.horizon, + impl_->config.action_dim, actions, error); +} + +const PIV3PolicyConfig & PIV3Policy::config() const { + if (impl_ == nullptr) { + throw std::runtime_error("StarVLA PI_v3 policy is not initialized"); + } + return impl_->config; +} + +const char * PIV3Policy::backend_name() const { + return impl_ != nullptr ? mode_name(impl_->mode) : "unknown"; +} + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/pi_v3_policy.h b/src/models/starvla/pi_v3_policy.h new file mode 100644 index 0000000..8177221 --- /dev/null +++ b/src/models/starvla/pi_v3_policy.h @@ -0,0 +1,98 @@ +#pragma once + +#include "models/starvla/normalization.h" + +#include +#include +#include +#include +#include + +namespace robotcpp::starvla { + +struct PIV3PolicyConfig { + std::string backbone_arch; + std::string bundle_uuid; + std::string text_filename; + std::string mmproj_filename; + + int qwen_hidden_dim = 0; + int qwen_input_embedding_dim = 0; + int qwen_layer_count = 0; + int qwen_vocab_size = 0; + std::string cot_template; + int image_count = 0; + std::vector image_names; + int image_processor_min_pixels = 0; + int image_processor_max_pixels = 0; + int image_patch_size = 0; + int image_spatial_merge_size = 0; + int image_min_token_count = 0; + int image_max_token_count = 0; + + int dit_width = 0; + int block_count = 0; + int projector_count = 0; + int attention_head_count = 0; + int attention_head_dim = 0; + int feed_forward_dim = 0; + int mlp_hidden_dim = 0; + int action_dim = 0; + int horizon = 0; + int future_token_count = 0; + int action_position_count = 0; + int no_state_sequence_length = 0; + int timestep_projection_dim = 0; + int num_timestep_buckets = 0; + int num_inference_timesteps = 0; + float ada_norm_epsilon = 0.0f; + float projector_norm_epsilon = 0.0f; + float euler_dt = 0.0f; + std::vector qwen_hidden_tuple_indices; + std::vector timestep_ids; + NormalizationConfig normalization; +}; + +class PIV3Policy { + public: + ~PIV3Policy(); + + PIV3Policy(const PIV3Policy &) = delete; + PIV3Policy & operator=(const PIV3Policy &) = delete; + + static std::unique_ptr load(const std::string & path, int n_threads, int verbosity, + std::string & error); + + // qwen_hidden_states is layer-major + // [qwen_layer_count, token_count, qwen_hidden_dim]. Each layer has the + // same full-chat token sequence. Non-zero mask entries participate in + // every cross-attention block. The released checkpoint has no raw-state + // runtime path. initial_noise is token-major [horizon, action_dim]. + bool evaluate(const float * qwen_hidden_states, size_t hidden_element_count, + const uint8_t * qwen_attention_mask, size_t mask_element_count, + const float * initial_noise, size_t noise_element_count, + std::vector & normalized_actions, std::string & error); + bool unnormalize(const std::vector & normalized_actions, const std::string & profile_key, + std::vector & actions, std::string & error) const; + + const PIV3PolicyConfig & config() const; + const char * backend_name() const; + + private: + struct Impl; + + explicit PIV3Policy(std::unique_ptr impl); + + bool evaluate_internal(const float * qwen_hidden_states, + size_t hidden_element_count, + const uint8_t * qwen_attention_mask, + size_t mask_element_count, + const float * initial_noise, + size_t noise_element_count, + std::vector & normalized_actions, + std::string & error); + + std::unique_ptr impl_; +}; + +} // namespace robotcpp::starvla diff --git a/tests/starvla/test_starvla_pi_v3_golden.py b/tests/starvla/test_starvla_pi_v3_golden.py new file mode 100644 index 0000000..4551bbd --- /dev/null +++ b/tests/starvla/test_starvla_pi_v3_golden.py @@ -0,0 +1,36 @@ +from __future__ import annotations + +import sys +import unittest +from pathlib import Path + + +TOOLS_DIR = Path(__file__).resolve().parents[2] / "tools" / "hf2gguf" / "starvla" +sys.path.insert(0, str(TOOLS_DIR)) + +from generate_starvla_pi_v3_golden import ( # noqa: E402 + CONDITIONING_TAP_NAMES, + expected_model_instruction, + expected_runtime_contract, +) + + +class PIv3ReferenceTest(unittest.TestCase): + def test_instruction_template(self) -> None: + config = { + "datasets": { + "vla_data": {"CoT_prompt": "Task: {instruction}"}, + } + } + self.assertEqual(expected_model_instruction(config, "grab block"), "Task: grab block") + + def test_action_oracle_contract(self) -> None: + contract = expected_runtime_contract() + self.assertEqual(contract["conditioning"]["hidden_tuple_indices"], list(range(1, 37))) + self.assertEqual(contract["conditioning"]["hidden_tap_names"], CONDITIONING_TAP_NAMES) + self.assertEqual(contract["timesteps"], [0, 250, 500, 750]) + self.assertEqual(contract["action_shape"], [16, 7]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tools/hf2gguf/starvla/generate_starvla_pi_v3_golden.py b/tools/hf2gguf/starvla/generate_starvla_pi_v3_golden.py new file mode 100644 index 0000000..1bb0b91 --- /dev/null +++ b/tools/hf2gguf/starvla/generate_starvla_pi_v3_golden.py @@ -0,0 +1,1831 @@ +#!/usr/bin/env python3 +"""Generate an auditable oracle from the pinned official StarVLA PI_v3 checkpoint. + +The exporter executes the pinned StarVLA QwenPI_v3 implementation and records +both Transformers' effective outer-model conditioning tuple and cloned raw +decoder-layer outputs. This distinction matters in Transformers 4.57: +DeepStack updates the first three recorded decoder outputs in place, while the +outer conditional model retains the raw final decoder output. +""" + +from __future__ import annotations + +import argparse +import contextlib +import datetime as dt +import gc +import hashlib +import importlib.metadata +import json +import math +import os +import platform +import random +import shutil +import subprocess +import sys +import tempfile +from pathlib import Path +from types import MethodType +from typing import Any, Iterable, Mapping, Sequence + +import numpy as np + + +TOOLS_DIR = Path(__file__).resolve().parent +if str(TOOLS_DIR) not in sys.path: + sys.path.insert(0, str(TOOLS_DIR)) + +from starvla_checkpoint import ( # noqa: E402 + DEFAULT_CATALOG, + StarVLAError, + get_variant, + load_catalog, + official_bundle_uuid, + resolve_effective_config, + sha256_file, + verify_catalog_files, + verify_checkpoint_file, +) + + +GOLDEN_SCHEMA_VERSION = 1 +SUPPORTED_VARIANT = "pi_v3" +GOLDEN_KIND = "starvla_pi_v3_official_python_oracle" +SEED = 0 +EXPECTED_TRANSFORMERS_VERSION = "4.57.0" +EXPECTED_TORCH_VERSION = "2.6.0" +EXPECTED_TORCHVISION_VERSION = "0.21.0" +EXPECTED_NUMPY_VERSION = "1.26.4" +EXPECTED_DIFFUSERS_VERSION = "0.37.1" +EXPECTED_TOKENIZERS_VERSION = "0.22.2" +EXPECTED_PILLOW_VERSION = "12.1.1" +EXPECTED_OMEGACONF_VERSION = "2.3.0" +EXPECTED_ACCELERATE_VERSION = "1.5.2" +EXPECTED_SAFETENSORS_VERSION = "0.7.0" +EXPECTED_QWEN3VL_MODELING_SHA256 = "dd63ed3b124232735b3dca1bfa28f9d6b0d3f7182afcb75dde8f3e724b2b22da" +EXPECTED_TRANSFORMERS_GENERIC_SHA256 = "b117ffb2e9d513def41ce596eb82057b8e2811c6edf29ffd0bb634979240ebed" +EXPECTED_QWEN3VL_PROCESSING_SHA256 = "efd8d64aaf608aad1ffb3e6d503d6a99e5227d007df95c1d9fa905d998cda4a9" +EXPECTED_QWEN2VL_IMAGE_PROCESSING_FAST_SHA256 = ( + "09bfa9b17df7c3f0c6159bc34008ee50f21d2472cd5bae7e5c21ba1ca13a423c" +) +EXPECTED_QWEN2VL_IMAGE_PROCESSING_SHA256 = ( + "7820a0fcca107e75605e08d9b774285ca2b0316f857bc0225c779794705ecf4f" +) +OFFICIAL_ENVIRONMENT_FREEZE = { + "path": "wandb/wandb/run-20260426_011111-enstjn5q/files/requirements.txt", + "size": 4354, + "sha256": "de6b505238663ea8a218620e8a4f99cbcfe1e6e09f347ab26f68fe434f3fb00e", +} +EXPECTED_ACTION_HORIZON = 16 +EXPECTED_ACTION_DIM = 7 +EXPECTED_LAYER_COUNT = 36 +EXPECTED_QWEN_HIDDEN_DIM = 2560 +EXPECTED_PROJECTED_HIDDEN_DIM = 1024 +EXPECTED_TIMESTEP_IDS = [0, 250, 500, 750] +EXPECTED_COT_TEMPLATE = ( + "Your task is {instruction}. To identify the key objects for your task. " + "Locate their bounding boxes in [x1,y1,x2,y2] format." +) +CONDITIONING_TAP_NAMES = ( + [f"deepstack_out-{index}" for index in range(3)] + + [f"l_out-{index}" for index in range(3, EXPECTED_LAYER_COUNT)] +) +RAW_TAP_NAMES = [f"l_out-{index}" for index in range(EXPECTED_LAYER_COUNT)] +FINAL_NORM_DIAGNOSTIC_NAME = "result_norm" +CONDITIONING_SEMANTICS = ( + "Transformers 4.57 outer conditional-model recorder references after in-place DeepStack, " + "then raw decoder outputs including the final layer" +) +PROJECTOR_AUTOCAST_CONTRACT = { + "autocast_device_type": "cuda", + "autocast_dtype": "bfloat16", + "layer_norm_input_dtype": "bfloat16", + "layer_norm_parameter_dtype": "float32", + "layer_norm_compute_dtype": "float32", + "layer_norm_output_dtype": "float32", + "linear_input_operand_dtype": "bfloat16", + "linear_weight_operand_dtype": "bfloat16", + "linear_bias_operand_dtype": "bfloat16", + "linear_bias_application": "cublaslt_epilogue_bias", + "linear_operand_rounding": "round_to_nearest_even", + "linear_per_split_accumulator_dtype": "float32", + "linear_split_partial_storage_dtype": "bfloat16", + "linear_split_reduction_scheme": "output_type", + "allow_bf16_reduced_precision_reduction_setting_affects_gemm_and_bias": False, + "linear_output_dtype": "bfloat16", + "saved_output_dtype": "float32", + "saved_output_transport": "exact_widen_of_bfloat16_value", + "layer_norm_validation": "all_36_outputs_bitwise_equal_explicit_fp32_reconstruction", + "linear_validation": "all_36_outputs_bitwise_equal_explicit_bf16_operand_reconstruction", + "fp32_linear_then_output_round_is_distinct": True, +} + + +def _canonical_json(value: Any) -> bytes: + return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode("utf-8") + + +def _sha256_bytes(value: bytes) -> str: + return hashlib.sha256(value).hexdigest() + + +def _array_sha256(value: np.ndarray) -> str: + array = np.ascontiguousarray(value) + header = _canonical_json({"dtype": array.dtype.str, "shape": list(array.shape)}) + digest = hashlib.sha256() + digest.update(header) + digest.update(b"\x00") + payload = memoryview(array).cast("B") + for start in range(0, len(payload), 16 * 1024 * 1024): + digest.update(payload[start : start + 16 * 1024 * 1024]) + return digest.hexdigest() + + +def _array_record(value: np.ndarray, *, source_dtype: str | None = None) -> dict[str, Any]: + array = np.ascontiguousarray(value) + record: dict[str, Any] = { + "dtype": array.dtype.str, + "shape": list(array.shape), + "sha256": _array_sha256(array), + } + if source_dtype is not None: + record["source_dtype"] = source_dtype + return record + + +def _distribution_version(name: str) -> str: + try: + return importlib.metadata.version(name) + except importlib.metadata.PackageNotFoundError: + return "missing" + + +def _base_version(version: str) -> str: + return version.split("+", 1)[0] + + +def validate_runtime_versions( + *, + torch_version: str, + torchvision_version: str, + transformers_version: str, + numpy_version: str, + diffusers_version: str, + tokenizers_version: str, + pillow_version: str, + omegaconf_version: str, + accelerate_version: str, + safetensors_version: str, +) -> None: + expected = { + "torch": EXPECTED_TORCH_VERSION, + "torchvision": EXPECTED_TORCHVISION_VERSION, + "transformers": EXPECTED_TRANSFORMERS_VERSION, + "numpy": EXPECTED_NUMPY_VERSION, + "diffusers": EXPECTED_DIFFUSERS_VERSION, + "tokenizers": EXPECTED_TOKENIZERS_VERSION, + "pillow": EXPECTED_PILLOW_VERSION, + "omegaconf": EXPECTED_OMEGACONF_VERSION, + "accelerate": EXPECTED_ACCELERATE_VERSION, + "safetensors": EXPECTED_SAFETENSORS_VERSION, + } + actual = { + "torch": _base_version(torch_version), + "torchvision": _base_version(torchvision_version), + "transformers": _base_version(transformers_version), + "numpy": _base_version(numpy_version), + "diffusers": _base_version(diffusers_version), + "tokenizers": _base_version(tokenizers_version), + "pillow": _base_version(pillow_version), + "omegaconf": _base_version(omegaconf_version), + "accelerate": _base_version(accelerate_version), + "safetensors": _base_version(safetensors_version), + } + mismatches = [ + f"{name}: expected {expected[name]}, got {actual[name]}" + for name in expected + if actual[name] != expected[name] + ] + if mismatches: + raise StarVLAError("official PI_v3 oracle runtime version mismatch: " + "; ".join(mismatches)) + + +def expected_model_instruction(config: Mapping[str, Any], task: str) -> str: + if not isinstance(task, str) or not task.strip(): + raise StarVLAError("task must be a non-empty string") + try: + vla_data = config["datasets"]["vla_data"] + except (KeyError, TypeError) as exc: + raise StarVLAError("checkpoint config has no datasets.vla_data object") from exc + if not isinstance(vla_data, Mapping): + raise StarVLAError("checkpoint config datasets.vla_data must be an object") + cot_prompt = vla_data.get("CoT_prompt") + if not isinstance(cot_prompt, str) or cot_prompt.count("{instruction}") != 1: + raise StarVLAError("official PI_v3 CoT_prompt must contain exactly one {instruction} placeholder") + return cot_prompt.replace("{instruction}", task) + + +def expected_runtime_contract() -> dict[str, Any]: + """Describe the reference inputs needed to reproduce the action oracle.""" + return { + "conditioning": { + "hidden_tuple_indices": list(range(1, 37)), + "hidden_tap_names": CONDITIONING_TAP_NAMES, + }, + "timesteps": EXPECTED_TIMESTEP_IDS, + "action_shape": [16, 7], + } + + +def _run_git(source_dir: Path, *arguments: str) -> str: + try: + result = subprocess.run( + ["git", "-C", str(source_dir), *arguments], + check=True, + capture_output=True, + text=True, + ) + except (OSError, subprocess.CalledProcessError) as exc: + raise StarVLAError(f"failed to inspect pinned StarVLA checkout {source_dir}: {exc}") from exc + return result.stdout.strip() + + +def verify_pinned_source_checkout(source_dir: Path, expected_revision: str) -> None: + source_dir = source_dir.resolve() + if not (source_dir / ".git").exists(): + raise StarVLAError(f"StarVLA source is not a Git checkout: {source_dir}") + actual_revision = _run_git(source_dir, "rev-parse", "HEAD") + if actual_revision != expected_revision: + raise StarVLAError( + f"StarVLA source revision mismatch: expected {expected_revision}, got {actual_revision}" + ) + changes = _run_git(source_dir, "status", "--porcelain=v1", "--untracked-files=all") + if changes: + raise StarVLAError(f"pinned StarVLA checkout has tracked or untracked changes:\n{changes}") + + +def _ensure_regular_file(path: Path, *, label: str) -> None: + if not path.is_file() or path.is_symlink(): + raise StarVLAError(f"{label} must be a regular, non-symlink file: {path}") + + +def validate_official_inputs( + *, + checkpoint_root: Path, + source_dir: Path, + catalog_path: Path = DEFAULT_CATALOG, +) -> dict[str, Any]: + catalog = load_catalog(catalog_path) + variant = get_variant(catalog, SUPPORTED_VARIANT) + qwen = catalog["shared_assets"]["qwen3_vl_4b_instruct"] + checkpoint_root = checkpoint_root.resolve() + policy_dir = checkpoint_root / "sources" / variant["directory"] / variant["revision"] + qwen_dir = checkpoint_root / "sources" / qwen["directory"] / qwen["revision"] + checkpoint = policy_dir / variant["checkpoint"]["path"] + + expected_source = (checkpoint_root / "source" / "starvla").resolve() + if source_dir.resolve() != expected_source: + raise StarVLAError( + f"StarVLA source must be the canonical checkout {expected_source}, got {source_dir.resolve()}" + ) + verify_pinned_source_checkout(source_dir, catalog["source_revisions"]["starvla"]) + verify_catalog_files(policy_dir, variant) + verify_catalog_files(qwen_dir, qwen) + _ensure_regular_file(checkpoint, label="official PI_v3 checkpoint") + incomplete_sidecar = Path(f"{checkpoint}.aria2") + if incomplete_sidecar.exists(): + raise StarVLAError( + f"official PI_v3 checkpoint download is incomplete ({incomplete_sidecar} exists); resume it first" + ) + verify_checkpoint_file(checkpoint, variant) + return { + "catalog": catalog, + "variant": variant, + "qwen": qwen, + "policy_dir": policy_dir, + "qwen_dir": qwen_dir, + "checkpoint": checkpoint, + "source_dir": source_dir.resolve(), + "catalog_path": catalog_path.resolve(), + } + + +def _require_isolated_python() -> None: + if not sys.flags.isolated: + raise StarVLAError( + "the PI_v3 oracle must run in isolated mode; invoke it as `python -I " + "tools/hf2gguf/starvla/generate_starvla_pi_v3_golden.py ...`" + ) + + +def _configure_determinism(torch: Any, *, seed: int, device: str) -> None: + os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" + os.environ["NO_ALBUMENTATIONS_UPDATE"] = "1" + if not device.startswith("cuda"): + raise StarVLAError("the official PI_v3 golden oracle requires a CUDA device") + if not torch.cuda.is_available(): + raise StarVLAError("CUDA is not available to PyTorch") + try: + cuda_device = torch.device(device) + except (RuntimeError, ValueError) as exc: + raise StarVLAError(f"invalid CUDA device {device!r}: {exc}") from exc + torch.cuda.set_device(0 if cuda_device.index is None else cuda_device.index) + if not torch.cuda.is_bf16_supported(): + raise StarVLAError(f"CUDA device {device!r} does not support bfloat16") + + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + torch.use_deterministic_algorithms(True) + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False + torch.backends.cudnn.allow_tf32 = False + torch.backends.cudnn.benchmark = False + + +def verify_transformers_qwen3vl_recorder_semantics(torch: Any, transformers: Any) -> dict[str, Any]: + """Execute the outer 4-layer Qwen3-VL model to gate 4.57 recorder semantics.""" + + try: + from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLConfig + from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLForConditionalGeneration + from transformers.models.qwen2_vl import ( + image_processing_qwen2_vl, + image_processing_qwen2_vl_fast, + ) + from transformers.models.qwen3_vl import modeling_qwen3_vl, processing_qwen3_vl + from transformers.utils import generic as transformers_generic + except ImportError as exc: + raise StarVLAError(f"Transformers lacks the pinned Qwen3-VL implementation: {exc}") from exc + modeling_path = Path(modeling_qwen3_vl.__file__).resolve() + actual_source_sha = sha256_file(modeling_path) + if actual_source_sha != EXPECTED_QWEN3VL_MODELING_SHA256: + raise StarVLAError( + "Transformers 4.57 Qwen3-VL implementation SHA256 mismatch: " + f"expected {EXPECTED_QWEN3VL_MODELING_SHA256}, got {actual_source_sha}" + ) + generic_path = Path(transformers_generic.__file__).resolve() + actual_generic_sha = sha256_file(generic_path) + if actual_generic_sha != EXPECTED_TRANSFORMERS_GENERIC_SHA256: + raise StarVLAError( + "Transformers 4.57 recorder implementation SHA256 mismatch: " + f"expected {EXPECTED_TRANSFORMERS_GENERIC_SHA256}, got {actual_generic_sha}" + ) + processing_path = Path(processing_qwen3_vl.__file__).resolve() + actual_processing_sha = sha256_file(processing_path) + if actual_processing_sha != EXPECTED_QWEN3VL_PROCESSING_SHA256: + raise StarVLAError( + "Transformers 4.57 Qwen3-VL processor implementation SHA256 mismatch: " + f"expected {EXPECTED_QWEN3VL_PROCESSING_SHA256}, got {actual_processing_sha}" + ) + image_processing_fast_path = Path(image_processing_qwen2_vl_fast.__file__).resolve() + actual_image_processing_fast_sha = sha256_file(image_processing_fast_path) + if actual_image_processing_fast_sha != EXPECTED_QWEN2VL_IMAGE_PROCESSING_FAST_SHA256: + raise StarVLAError( + "Transformers 4.57 Qwen2-VL fast image processor SHA256 mismatch: " + f"expected {EXPECTED_QWEN2VL_IMAGE_PROCESSING_FAST_SHA256}, " + f"got {actual_image_processing_fast_sha}" + ) + image_processing_path = Path(image_processing_qwen2_vl.__file__).resolve() + actual_image_processing_sha = sha256_file(image_processing_path) + if actual_image_processing_sha != EXPECTED_QWEN2VL_IMAGE_PROCESSING_SHA256: + raise StarVLAError( + "Transformers 4.57 Qwen2-VL smart-resize implementation SHA256 mismatch: " + f"expected {EXPECTED_QWEN2VL_IMAGE_PROCESSING_SHA256}, " + f"got {actual_image_processing_sha}" + ) + + config = Qwen3VLConfig( + text_config={ + "vocab_size": 32, + "hidden_size": 16, + "intermediate_size": 32, + "num_hidden_layers": 4, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 8, + "max_position_embeddings": 32, + "use_cache": False, + "rope_scaling": { + "rope_type": "default", + "mrope_section": [2, 2, 4], + "mrope_interleaved": True, + }, + }, + vision_config={ + "depth": 1, + "hidden_size": 16, + "intermediate_size": 32, + "num_heads": 2, + "in_channels": 3, + "patch_size": 2, + "spatial_merge_size": 1, + "temporal_patch_size": 1, + "out_hidden_size": 16, + "num_position_embeddings": 16, + "deepstack_visual_indexes": [], + }, + image_token_id=2, + video_token_id=3, + vision_start_token_id=1, + vision_end_token_id=4, + ) + model = Qwen3VLForConditionalGeneration(config).cpu().eval() + raw: dict[int, Any] = {} + final_norm: dict[str, Any] = {} + inner_hidden: dict[str, Any] = {} + handles = [ + layer.register_forward_hook( + lambda _module, _inputs, output, index=index: raw.__setitem__( + index, output.detach().clone() + ) + ) + for index, layer in enumerate(model.model.language_model.layers) + ] + handles.append( + model.model.language_model.norm.register_forward_hook( + lambda _module, _inputs, output: final_norm.__setitem__( + "value", output.detach().clone() + ) + ) + ) + + def capture_inner(_module: Any, _inputs: Any, output: Any) -> None: + hidden_states = getattr(output, "hidden_states", None) + if hidden_states is not None: + inner_hidden["value"] = tuple(value.detach().clone() for value in hidden_states) + + handles.append(model.model.language_model.register_forward_hook(capture_inner)) + + def fake_get_image_features(_model: Any, pixel_values: Any, image_grid_thw: Any = None): + dtype = model.model.language_model.embed_tokens.weight.dtype + image_embed = torch.arange(16, dtype=dtype).reshape(1, 16) / 100.0 + deepstack = [ + torch.full((1, 16), float(index + 1), dtype=dtype) + for index in range(3) + ] + return [image_embed], deepstack + + model.model.get_image_features = MethodType(fake_get_image_features, model.model) + try: + input_ids = torch.tensor([[1, 2, 4, 5, 6]], dtype=torch.long) + visual_mask = input_ids == config.image_token_id + with torch.no_grad(): + output = model( + input_ids=input_ids, + attention_mask=torch.ones_like(input_ids), + pixel_values=torch.zeros(1), + image_grid_thw=torch.tensor([[1, 1, 1]], dtype=torch.long), + output_hidden_states=True, + use_cache=False, + logits_to_keep=1, + ) + finally: + for handle in handles: + handle.remove() + hidden = output.hidden_states + if hidden is None or len(hidden) != 5: + raise StarVLAError( + "Transformers 4.57 outer recorder probe did not return input + four layer states" + ) + for index in range(3): + delta = hidden[index + 1][visual_mask] - raw[index][visual_mask] + expected = torch.full_like(delta, float(index + 1)) + if not torch.allclose(delta, expected, rtol=0.0, atol=5e-7): + raise StarVLAError( + f"Transformers 4.57 recorder probe did not retain DeepStack's in-place layer {index} update" + ) + if not torch.equal(hidden[index + 1][~visual_mask], raw[index][~visual_mask]): + raise StarVLAError( + f"Transformers 4.57 recorder probe unexpectedly changed non-visual layer {index} tokens" + ) + if not torch.equal(hidden[-1], raw[3]): + raise StarVLAError( + "Transformers 4.57 outer recorder probe did not retain the raw final decoder output" + ) + if "value" not in final_norm: + raise StarVLAError("Transformers 4.57 recorder probe did not capture final RMSNorm") + if torch.equal(hidden[-1], final_norm["value"]): + raise StarVLAError( + "Transformers 4.57 outer recorder probe unexpectedly exposed result_norm as conditioning" + ) + inner = inner_hidden.get("value") + if inner is None or len(inner) != 5 or not torch.equal(inner[-1], final_norm["value"]): + raise StarVLAError( + "Transformers 4.57 inner recorder probe did not expose result_norm for diagnostics" + ) + return { + "modeling_qwen3_vl_path": str(modeling_path), + "modeling_qwen3_vl_sha256": actual_source_sha, + "transformers_generic_path": str(generic_path), + "transformers_generic_sha256": actual_generic_sha, + "processing_qwen3_vl_path": str(processing_path), + "processing_qwen3_vl_sha256": actual_processing_sha, + "image_processing_qwen2_vl_fast_path": str(image_processing_fast_path), + "image_processing_qwen2_vl_fast_sha256": actual_image_processing_fast_sha, + "image_processing_qwen2_vl_path": str(image_processing_path), + "image_processing_qwen2_vl_sha256": actual_image_processing_sha, + "model_class": "Qwen3VLForConditionalGeneration", + "observed_order": [ + "deepstack_out-0", + "deepstack_out-1", + "deepstack_out-2", + "l_out-3", + ], + "inner_terminal": "result_norm", + "outer_terminal": "l_out-3", + "mechanism": ( + "outer_recorder_keeps_raw_final_decoder_output_while_first_three_layer_references_receive_" + "in_place_deepstack_updates" + ), + } + + +@contextlib.contextmanager +def _config_only_qwen_bootstrap(torch: Any, transformers: Any, qwen_dir: Path): + model_class = transformers.Qwen3VLForConditionalGeneration + had_local_override = "from_pretrained" in model_class.__dict__ + original_local_override = model_class.__dict__.get("from_pretrained") + + def from_config_only(model_id: str | os.PathLike[str], **kwargs: Any): + actual = Path(model_id).resolve() + if actual != qwen_dir.resolve(): + raise StarVLAError(f"official wrapper requested unexpected Qwen source: {actual}") + if kwargs.get("dtype") not in (None, torch.bfloat16): + raise StarVLAError(f"unexpected Qwen bootstrap dtype: {kwargs.get('dtype')!r}") + config = transformers.AutoConfig.from_pretrained( + actual, + local_files_only=True, + trust_remote_code=False, + ) + if getattr(config, "model_type", None) != "qwen3_vl": + raise StarVLAError(f"unexpected pinned Qwen model_type: {getattr(config, 'model_type', None)!r}") + previous_dtype = torch.get_default_dtype() + try: + torch.set_default_dtype(torch.bfloat16) + with transformers.modeling_utils.no_init_weights(): + model = model_class(config) + finally: + torch.set_default_dtype(previous_dtype) + return model + + model_class.from_pretrained = staticmethod(from_config_only) + try: + yield + finally: + if had_local_override: + model_class.from_pretrained = original_local_override + else: + delattr(model_class, "from_pretrained") + + +@contextlib.contextmanager +def _official_qwen_model_alias(qwen_dir: Path): + """Expose the pinned local assets under StarVLA's case-sensitive dispatch name.""" + + qwen_dir = qwen_dir.resolve() + if not qwen_dir.is_dir(): + raise StarVLAError(f"pinned Qwen asset directory does not exist: {qwen_dir}") + with tempfile.TemporaryDirectory(prefix="starvla-qwen-alias-") as temporary: + alias = Path(temporary) / "Qwen3-VL-4B-Instruct" + alias.symlink_to(qwen_dir, target_is_directory=True) + if not alias.is_dir() or alias.resolve() != qwen_dir: + raise StarVLAError(f"temporary Qwen alias did not resolve to the pinned model: {alias}") + yield alias + + +def _assert_module_origin(module: Any, source_dir: Path) -> None: + module_path = Path(module.__file__).resolve() + try: + module_path.relative_to(source_dir.resolve()) + except ValueError as exc: + raise StarVLAError(f"imported StarVLA module is outside the pinned checkout: {module_path}") from exc + + +def verify_official_framework_import(paths: Mapping[str, Any]) -> dict[str, str]: + """Smoke-import the policy and normalizer from the already verified checkout.""" + source_dir = Path(paths["source_dir"]) + os.environ["NO_ALBUMENTATIONS_UPDATE"] = "1" + if any(name == "starVLA" or name.startswith("starVLA.") for name in sys.modules): + raise StarVLAError("starVLA was imported before pinned-source verification") + sys.path.insert(0, str(source_dir)) + try: + from deployment.model_server import policy_norm_processor + from starVLA.model.framework import base_framework, share_tools + from starVLA.model.framework.VLM4A import QwenPI_v3 + + modules = { + "base_framework": base_framework, + "share_tools": share_tools, + "qwen_pi_v3": QwenPI_v3, + "policy_norm_processor": policy_norm_processor, + } + for module in modules.values(): + _assert_module_origin(module, source_dir) + return {name: str(Path(module.__file__).resolve()) for name, module in modules.items()} + except ImportError as exc: + raise StarVLAError(f"failed to import the pinned official PI_v3 framework: {exc}") from exc + finally: + if sys.path and sys.path[0] == str(source_dir): + del sys.path[0] + + +def _validate_effective_config(config: Mapping[str, Any]) -> None: + try: + framework = config["framework"] + action = framework["action_model"] + diffusion = action["diffusion_model_cfg"] + vla_data = config["datasets"]["vla_data"] + except (KeyError, TypeError) as exc: + raise StarVLAError("effective PI_v3 config is missing required objects") from exc + actual = { + "framework": framework.get("name"), + "action_model_type": action.get("action_model_type"), + "action_horizon": action.get("action_horizon"), + "action_dim": action.get("action_dim"), + "state_dim": action.get("state_dim"), + "num_inference_timesteps": action.get("num_inference_timesteps"), + "num_timestep_buckets": action.get("num_timestep_buckets"), + "dit_width": diffusion.get("input_embedding_dim"), + "dit_layers": diffusion.get("num_layers"), + "interleave_self_attention": diffusion.get("interleave_self_attention"), + "use_canonical_forward": diffusion.get("use_canonical_forward"), + "image_size": vla_data.get("image_size"), + "data_mix": vla_data.get("data_mix"), + } + expected = { + "framework": "QwenPI_v3", + "action_model_type": "LayerwiseFM", + "action_horizon": 16, + "action_dim": 7, + "state_dim": 7, + "num_inference_timesteps": 4, + "num_timestep_buckets": 1000, + "dit_width": 1024, + "dit_layers": 36, + "interleave_self_attention": False, + "use_canonical_forward": True, + "image_size": [224, 224], + "data_mix": "bridge_rt_1", + } + if actual != expected: + raise StarVLAError(f"unexpected effective official PI_v3 config: {actual}") + if vla_data.get("CoT_prompt") != EXPECTED_COT_TEMPLATE: + raise StarVLAError(f"unexpected official PI_v3 CoT prompt: {vla_data.get('CoT_prompt')!r}") + expected_model_instruction(config, "contract probe") + + +def load_official_framework(paths: Mapping[str, Any], *, device: str) -> tuple[Any, dict[str, Any]]: + import torch + import transformers + + source_dir = Path(paths["source_dir"]) + if any(name == "starVLA" or name.startswith("starVLA.") for name in sys.modules): + raise StarVLAError("starVLA was imported before pinned-source verification") + sys.path.insert(0, str(source_dir)) + try: + from starVLA.model.framework import base_framework, share_tools + from starVLA.model.framework.VLM4A import QwenPI_v3 + + _assert_module_origin(base_framework, source_dir) + _assert_module_origin(share_tools, source_dir) + _assert_module_origin(QwenPI_v3, source_dir) + config = resolve_effective_config(Path(paths["policy_dir"]), SUPPORTED_VARIANT) + _validate_effective_config(config) + qwen_dir = Path(paths["qwen_dir"]).resolve() + with _official_qwen_model_alias(qwen_dir) as qwen_alias: + config = base_framework.merge_config_overrides( + config, + [ + f"framework.qwenvl.base_vlm={qwen_alias}", + "framework.qwenvl.attn_implementation=sdpa", + ], + ) + configured_qwen = Path(config["framework"]["qwenvl"]["base_vlm"]) + if "Qwen3-VL" not in str(configured_qwen) or configured_qwen.resolve() != qwen_dir: + raise StarVLAError( + "effective PI_v3 Qwen source does not preserve official dispatch and pinned assets" + ) + cfg = share_tools.dict_to_namespace(config) + cfg.trainer.pretrained_checkpoint = None + with _config_only_qwen_bootstrap(torch, transformers, qwen_dir): + framework = QwenPI_v3.Qwen_PI_v3(cfg) + + try: + state_dict = torch.load( + paths["checkpoint"], map_location="cpu", mmap=True, weights_only=True + ) + except TypeError: + state_dict = torch.load(paths["checkpoint"], map_location="cpu", weights_only=True) + if not isinstance(state_dict, Mapping) or not state_dict: + raise StarVLAError("official checkpoint did not contain a non-empty state_dict") + framework.load_state_dict(state_dict, strict=True) + del state_dict + gc.collect() + + if type(framework).__name__ != "Qwen_PI_v3": + raise StarVLAError(f"unexpected official framework class: {type(framework).__name__}") + action_model = framework.action_model + if int(framework.action_horizon) != EXPECTED_ACTION_HORIZON: + raise StarVLAError(f"unexpected official PI_v3 action horizon: {framework.action_horizon}") + if int(action_model.action_dim) != EXPECTED_ACTION_DIM: + raise StarVLAError(f"unexpected official PI_v3 action dimension: {action_model.action_dim}") + if len(framework.project_layers) != EXPECTED_LAYER_COUNT: + raise StarVLAError("official PI_v3 projector count is not 36") + if len(action_model.model.transformer_blocks) != EXPECTED_LAYER_COUNT: + raise StarVLAError("official PI_v3 DiT block count is not 36") + + qwen_dtypes = {parameter.dtype for parameter in framework.qwen_vl_interface.parameters()} + policy_dtypes = {parameter.dtype for parameter in action_model.parameters()} + projector_dtypes = {parameter.dtype for parameter in framework.project_layers.parameters()} + if qwen_dtypes != {torch.bfloat16}: + raise StarVLAError(f"unexpected Qwen parameter dtypes after strict load: {qwen_dtypes}") + if policy_dtypes != {torch.float32} or projector_dtypes != {torch.float32}: + raise StarVLAError( + "official PI_v3 FP32 policy/projector compatibility baseline changed: " + f"policy={policy_dtypes}, projectors={projector_dtypes}" + ) + return framework.to(device).eval(), config + finally: + if sys.path and sys.path[0] == str(source_dir): + del sys.path[0] + + +def _tensor_to_array(tensor: Any) -> tuple[np.ndarray, str]: + source_dtype = str(tensor.dtype).removeprefix("torch.") + value = tensor.detach().cpu().contiguous() + if source_dtype == "bfloat16": + value = value.float() + return np.ascontiguousarray(value.numpy()), source_dtype + + +def _require_tensor_equal(torch: Any, actual: Any, expected: Any, label: str) -> None: + if actual.shape != expected.shape or actual.dtype != expected.dtype or not torch.equal(actual, expected): + raise StarVLAError(f"official PI_v3 instrumentation mismatch for {label}") + + +def _projector_linear_bf16_operands( + torch: Any, + value: Any, + weight: Any, + bias: Any, +) -> Any: + """Replay CUDA autocast's Linear policy with explicit BF16 operands.""" + + with torch.autocast(value.device.type, enabled=False): + output = torch.nn.functional.linear( + value.to(dtype=torch.bfloat16), + weight.to(dtype=torch.bfloat16), + None if bias is None else bias.to(dtype=torch.bfloat16), + ) + if output.dtype != torch.bfloat16: + raise StarVLAError("explicit PI_v3 projector BF16 replay did not produce BF16") + return output + + +def _projector_linear_fp32_then_bf16( + torch: Any, + value: Any, + weight: Any, + bias: Any, +) -> Any: + """Represent the rejected FP32-Linear-then-BF16-round interpretation.""" + + with torch.autocast(value.device.type, enabled=False): + return torch.nn.functional.linear( + value.to(dtype=torch.float32), + weight.to(dtype=torch.float32), + None if bias is None else bias.to(dtype=torch.float32), + ).to(dtype=torch.bfloat16) + + +def run_official_forward( + framework: Any, + *, + images: Sequence[Any], + task: str, + seed: int = SEED, +) -> dict[str, Any]: + """Execute Qwen_PI_v3.predict_action and capture every parity boundary.""" + + import torch + + if torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction: + raise StarVLAError( + "PI_v3 official forward requires BF16 GEMM reduced-precision reduction to be disabled" + ) + captures: dict[str, Any] = {} + qwen = framework.qwen_vl_interface + action_model = framework.action_model + language_model = qwen.model.model.language_model + original_build = qwen.build_qwenvl_inputs + original_project = framework._project_vl_hidden_for_action + original_policy = action_model.predict_action + original_action_encoder = action_model.action_encoder.forward + original_dit = action_model.model.forward + raw_layer_outputs: dict[int, Any] = {} + deepstack_outputs: dict[int, Any] = {} + final_norm: dict[str, Any] = {} + projector_norm_inputs: dict[int, Any] = {} + projector_norm_outputs: dict[int, Any] = {} + projector_linear_inputs: dict[int, Any] = {} + projector_linear_outputs: dict[int, Any] = {} + projector_autocast_states: dict[tuple[int, str], tuple[bool, Any]] = {} + handles = [] + + def record_projector_tensor( + storage: dict[int, Any], + index: int, + value: Any, + label: str, + ) -> None: + if index in storage or not isinstance(value, torch.Tensor): + raise StarVLAError(f"official PI_v3 projector hook mismatch for {label} {index}") + storage[index] = value.detach() + + def record_projector_autocast(index: int, stage: str) -> None: + key = (index, stage) + if key in projector_autocast_states: + raise StarVLAError( + f"official PI_v3 projector autocast hook ran twice for {stage} {index}" + ) + projector_autocast_states[key] = ( + torch.is_autocast_enabled("cuda"), + torch.get_autocast_dtype("cuda"), + ) + + def capture_build(*args: Any, **kwargs: Any): + if "qwen_inputs" in captures: + raise StarVLAError("official PI_v3 preprocessing ran more than once") + batch_images = kwargs.get("images", args[0] if args else None) + instructions = kwargs.get("instructions", args[1] if len(args) > 1 else None) + captures["processed_images"] = list(batch_images[0]) + captures["framework_instructions"] = list(instructions) + result = original_build(*args, **kwargs) + captures["qwen_inputs"] = { + key: value.detach() for key, value in result.items() if isinstance(value, torch.Tensor) + } + return result + + def capture_project(hidden_states: Sequence[Any]): + if "project_input_taps" in captures: + raise StarVLAError("official PI_v3 projector bridge ran more than once") + captures["project_input_taps"] = [value.detach() for value in hidden_states] + projector_handles = [] + for index, projector in enumerate(framework.project_layers): + if ( + not isinstance(projector, torch.nn.Sequential) + or len(projector) != 2 + or not isinstance(projector[0], torch.nn.LayerNorm) + or not isinstance(projector[1], torch.nn.Linear) + ): + raise StarVLAError( + f"official PI_v3 projector {index} is no longer LayerNorm then Linear" + ) + + def capture_norm_input(_module: Any, inputs: Any, *, index: int = index) -> None: + if len(inputs) != 1: + raise StarVLAError(f"official PI_v3 projector norm {index} input arity changed") + record_projector_autocast(index, "layer_norm") + record_projector_tensor( + projector_norm_inputs, index, inputs[0], "LayerNorm input" + ) + + def capture_norm_output( + _module: Any, _inputs: Any, output: Any, *, index: int = index + ) -> None: + record_projector_tensor( + projector_norm_outputs, index, output, "LayerNorm output" + ) + + def capture_linear_input(_module: Any, inputs: Any, *, index: int = index) -> None: + if len(inputs) != 1: + raise StarVLAError(f"official PI_v3 projector linear {index} input arity changed") + record_projector_autocast(index, "linear") + record_projector_tensor( + projector_linear_inputs, index, inputs[0], "Linear logical input" + ) + + def capture_linear_output( + _module: Any, _inputs: Any, output: Any, *, index: int = index + ) -> None: + record_projector_tensor( + projector_linear_outputs, index, output, "Linear output" + ) + + projector_handles.extend( + [ + projector[0].register_forward_pre_hook(capture_norm_input), + projector[0].register_forward_hook(capture_norm_output), + projector[1].register_forward_pre_hook(capture_linear_input), + projector[1].register_forward_hook(capture_linear_output), + ] + ) + try: + projected = original_project(hidden_states) + finally: + for handle in projector_handles: + handle.remove() + captures["projected_hidden_taps"] = [value.detach() for value in projected] + return projected + + def capture_action_encoder(actions: Any, timesteps: Any): + return original_action_encoder(actions.to(dtype=torch.float32), timesteps) + + def capture_dit(*args: Any, **kwargs: Any): + conditioning = kwargs.get("encoder_hidden_states") + if not isinstance(conditioning, (list, tuple)): + raise StarVLAError("official PI_v3 DiT did not receive layer-wise conditioning") + kwargs["encoder_hidden_states"] = [value.to(dtype=torch.float32) for value in conditioning] + timestep = kwargs.get("timestep") + if timestep is None or timestep.numel() != 1: + raise StarVLAError("official PI_v3 DiT timestep shape changed") + captures.setdefault("timestep_ids", []).append(int(timestep.item())) + return original_dit(*args, **kwargs) + + def capture_policy(*args: Any, **kwargs: Any): + if "policy_input_taps" in captures: + raise StarVLAError("official PI_v3 policy sampler ran more than once") + policy_hidden = args[0] if args else kwargs.get("vl_embs_list") + if not isinstance(policy_hidden, (list, tuple)): + raise StarVLAError("official PI_v3 policy did not receive layer-wise hidden states") + captures["policy_input_taps"] = [value.detach() for value in policy_hidden] + original_randn = torch.randn + + def capture_randn(*randn_args: Any, **randn_kwargs: Any): + value = original_randn(*randn_args, **randn_kwargs) + if "initial_noise" in captures: + raise StarVLAError("official PI_v3 policy sampled noise more than once") + captures["initial_noise"] = value.detach().clone() + return value + + torch.randn = capture_randn + try: + output = original_policy(*args, **kwargs) + finally: + torch.randn = original_randn + captures["raw_policy"] = output.detach() + return output + + def capture_qwen_hidden(_module: Any, _inputs: Any, output: Any): + if "qwen_hidden_tuple" in captures: + raise StarVLAError("official PI_v3 outer Qwen model ran more than once") + hidden_states = getattr(output, "hidden_states", None) + if hidden_states is None or len(hidden_states) != EXPECTED_LAYER_COUNT + 1: + raise StarVLAError( + "official Transformers Qwen output did not contain input + 36 hidden states" + ) + captures["qwen_hidden_tuple"] = [value.detach() for value in hidden_states] + captures["conditioning_taps"] = [value.detach() for value in hidden_states[-36:]] + + for index, layer in enumerate(language_model.layers): + handles.append( + layer.register_forward_hook( + lambda _module, _inputs, output, index=index: raw_layer_outputs.__setitem__( + index, output.detach().clone() + ) + ) + ) + if index in (1, 2, 3): + handles.append( + layer.register_forward_pre_hook( + lambda _module, inputs, index=index: deepstack_outputs.__setitem__( + index - 1, inputs[0].detach().clone() + ) + ) + ) + handles.append( + language_model.norm.register_forward_hook( + lambda _module, _inputs, output: final_norm.__setitem__("value", output.detach().clone()) + ) + ) + handles.append(qwen.model.register_forward_hook(capture_qwen_hidden)) + qwen.build_qwenvl_inputs = capture_build + framework._project_vl_hidden_for_action = capture_project + action_model.predict_action = capture_policy + action_model.action_encoder.forward = capture_action_encoder + action_model.model.forward = capture_dit + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + try: + result = framework.predict_action(examples=[{"image": list(images), "lang": task}]) + finally: + for handle in handles: + handle.remove() + qwen.build_qwenvl_inputs = original_build + framework._project_vl_hidden_for_action = original_project + action_model.predict_action = original_policy + action_model.action_encoder.forward = original_action_encoder + action_model.model.forward = original_dit + + required = { + "processed_images", + "framework_instructions", + "qwen_inputs", + "qwen_hidden_tuple", + "conditioning_taps", + "project_input_taps", + "projected_hidden_taps", + "policy_input_taps", + "initial_noise", + "raw_policy", + "timestep_ids", + } + missing = sorted(required - set(captures)) + if missing: + raise StarVLAError(f"official PI_v3 instrumentation did not capture: {missing}") + if set(raw_layer_outputs) != set(range(EXPECTED_LAYER_COUNT)): + raise StarVLAError("official PI_v3 instrumentation missed raw Qwen decoder outputs") + if set(deepstack_outputs) != {0, 1, 2} or "value" not in final_norm: + raise StarVLAError("official PI_v3 instrumentation missed DeepStack/result_norm outputs") + for name in ( + "conditioning_taps", + "project_input_taps", + "projected_hidden_taps", + "policy_input_taps", + ): + if len(captures[name]) != EXPECTED_LAYER_COUNT: + raise StarVLAError(f"official PI_v3 {name} count is not 36") + + conditioning = captures["conditioning_taps"] + for index in range(3): + _require_tensor_equal(torch, conditioning[index], deepstack_outputs[index], CONDITIONING_TAP_NAMES[index]) + if torch.equal(conditioning[index], raw_layer_outputs[index]): + raise StarVLAError(f"DeepStack layer {index} did not change any recorded hidden-state value") + for index in range(3, EXPECTED_LAYER_COUNT): + _require_tensor_equal(torch, conditioning[index], raw_layer_outputs[index], f"l_out-{index}") + if torch.equal(conditioning[-1], final_norm["value"]): + raise StarVLAError("official outer Qwen conditioning unexpectedly ends at result_norm") + for actual, expected, name in zip(captures["project_input_taps"], conditioning, CONDITIONING_TAP_NAMES): + _require_tensor_equal(torch, actual, expected, f"projector input {name}") + for actual, expected in zip(captures["policy_input_taps"], captures["projected_hidden_taps"]): + _require_tensor_equal(torch, actual, expected, "projector-to-policy BF16 boundary") + + expected_projectors = set(range(EXPECTED_LAYER_COUNT)) + if any( + set(storage) != expected_projectors + for storage in ( + projector_norm_inputs, + projector_norm_outputs, + projector_linear_inputs, + projector_linear_outputs, + ) + ): + raise StarVLAError("official PI_v3 instrumentation missed a projector numeric boundary") + expected_autocast_keys = { + (index, stage) + for index in range(EXPECTED_LAYER_COUNT) + for stage in ("layer_norm", "linear") + } + if set(projector_autocast_states) != expected_autocast_keys: + raise StarVLAError("official PI_v3 instrumentation missed a projector autocast state") + + fp32_interpretation_is_distinct = False + with torch.inference_mode(): + for index, projector in enumerate(framework.project_layers): + norm_input = projector_norm_inputs[index] + norm_output = projector_norm_outputs[index] + linear_input = projector_linear_inputs[index] + linear_output = projector_linear_outputs[index] + projected_output = captures["projected_hidden_taps"][index] + if projector_autocast_states[(index, "layer_norm")] != (True, torch.bfloat16): + raise StarVLAError( + f"official PI_v3 projector {index} LayerNorm did not run under CUDA BF16 autocast" + ) + if projector_autocast_states[(index, "linear")] != (True, torch.bfloat16): + raise StarVLAError( + f"official PI_v3 projector {index} Linear did not run under CUDA BF16 autocast" + ) + if norm_input.dtype != torch.bfloat16: + raise StarVLAError(f"official PI_v3 projector {index} LayerNorm input is not BF16") + if norm_output.dtype != torch.float32 or linear_input.dtype != torch.float32: + raise StarVLAError( + f"official PI_v3 projector {index} LayerNorm did not expose an FP32 output" + ) + if linear_output.dtype != torch.bfloat16 or projected_output.dtype != torch.bfloat16: + raise StarVLAError(f"official PI_v3 projector {index} Linear output is not BF16") + _require_tensor_equal( + torch, norm_input, captures["project_input_taps"][index], f"projector {index} norm input" + ) + _require_tensor_equal( + torch, linear_input, norm_output, f"projector {index} norm-to-linear input" + ) + _require_tensor_equal( + torch, linear_output, projected_output, f"projector {index} linear output" + ) + + norm = projector[0] + linear = projector[1] + if ( + norm.weight is None + or norm.bias is None + or norm.weight.dtype != torch.float32 + or norm.bias.dtype != torch.float32 + or linear.weight.dtype != torch.float32 + or linear.bias is None + or linear.bias.dtype != torch.float32 + ): + raise StarVLAError( + f"official PI_v3 projector {index} FP32 parameter boundary changed" + ) + with torch.autocast(norm_input.device.type, enabled=False): + explicit_norm = torch.nn.functional.layer_norm( + norm_input.to(dtype=torch.float32), + norm.normalized_shape, + norm.weight, + norm.bias, + norm.eps, + ) + _require_tensor_equal( + torch, + explicit_norm, + norm_output, + f"projector {index} explicit FP32 LayerNorm reconstruction", + ) + explicit_bf16 = _projector_linear_bf16_operands( + torch, linear_input, linear.weight.detach(), linear.bias.detach() + ) + _require_tensor_equal( + torch, + explicit_bf16, + projected_output, + f"projector {index} explicit BF16 operand reconstruction", + ) + if not fp32_interpretation_is_distinct: + fp32_then_bf16 = _projector_linear_fp32_then_bf16( + torch, linear_input, linear.weight.detach(), linear.bias.detach() + ) + fp32_interpretation_is_distinct = not torch.equal( + fp32_then_bf16, projected_output + ) + if not fp32_interpretation_is_distinct: + raise StarVLAError( + "official PI_v3 projector sample does not distinguish BF16 operands from " + "FP32 Linear followed by BF16 output rounding" + ) + captures["projector_autocast_contract"] = dict(PROJECTOR_AUTOCAST_CONTRACT) + + captures["raw_qwen_taps"] = [ + raw_layer_outputs[index] for index in range(EXPECTED_LAYER_COUNT) + ] + captures["result_norm_diagnostic"] = final_norm["value"] + if captures["timestep_ids"] != EXPECTED_TIMESTEP_IDS: + raise StarVLAError( + f"official PI_v3 timestep order changed: {captures['timestep_ids']}" + ) + for name in ( + "conditioning_taps", + "raw_qwen_taps", + "result_norm_diagnostic", + "projected_hidden_taps", + "initial_noise", + ): + values = captures[name] if isinstance(captures[name], list) else [captures[name]] + if any(value.dtype != torch.bfloat16 for value in values): + raise StarVLAError(f"official PI_v3 {name} boundary is no longer BF16") + expected_noise_shape = (1, EXPECTED_ACTION_HORIZON, EXPECTED_ACTION_DIM) + if tuple(captures["initial_noise"].shape) != expected_noise_shape: + raise StarVLAError( + f"official PI_v3 initial noise shape mismatch: {tuple(captures['initial_noise'].shape)}" + ) + if captures["raw_policy"].dtype != torch.float32: + raise StarVLAError(f"official PI_v3 policy output is not FP32: {captures['raw_policy'].dtype}") + normalized = np.asarray(result.get("normalized_actions")) + raw_policy, _ = _tensor_to_array(captures["raw_policy"]) + if normalized.shape != expected_noise_shape or not np.array_equal(normalized, raw_policy): + raise StarVLAError("official normalized_actions differ from the captured PI_v3 policy output") + if not np.isfinite(normalized).all(): + raise StarVLAError("official PI_v3 policy produced NaN or infinite actions") + captures["normalized_actions"] = np.ascontiguousarray(normalized, dtype=np.float32) + return captures + + +def _image_pixel_sha256(image: Any) -> str: + header = _canonical_json({"mode": image.mode, "size": list(image.size)}) + return _sha256_bytes(header + b"\x00" + image.tobytes()) + + +def _image_record(path: Path, image: Any) -> dict[str, Any]: + return { + "source_path": str(path.resolve()), + "source_size": path.stat().st_size, + "source_sha256": sha256_file(path), + "decoded_mode": image.mode, + "decoded_size": list(image.size), + "decoded_pixel_sha256": _image_pixel_sha256(image), + } + + +def _processed_image_records(images: Sequence[Any]) -> list[dict[str, Any]]: + return [ + { + "index": index, + "mode": image.mode, + "size": list(image.size), + "pixel_sha256": _image_pixel_sha256(image), + } + for index, image in enumerate(images) + ] + + +def _render_model_prompt(framework: Any, processed_images: Sequence[Any], instruction: str) -> str: + messages = [ + { + "role": "user", + "content": [ + *({"type": "image", "image": image} for image in processed_images), + {"type": "text", "text": instruction}, + ], + } + ] + rendered = framework.qwen_vl_interface.processor.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True, + ) + if not isinstance(rendered, str): + raise StarVLAError(f"official processor returned a non-string prompt: {type(rendered)}") + return rendered + + +def _processor_patch_contract(framework: Any) -> dict[str, int]: + vision_config = framework.qwen_vl_interface.model.config.vision_config + image_processor = framework.qwen_vl_interface.processor.image_processor + + def positive_int(value: Any, label: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise StarVLAError(f"official Qwen processor {label} is not a positive integer: {value!r}") + return value + + channel_count = positive_int(vision_config.in_channels, "vision in_channels") + vision_patch_size = positive_int(vision_config.patch_size, "vision patch_size") + vision_temporal_patch_size = positive_int( + vision_config.temporal_patch_size, + "vision temporal_patch_size", + ) + processor_patch_size = positive_int(image_processor.patch_size, "image processor patch_size") + processor_temporal_patch_size = positive_int( + image_processor.temporal_patch_size, + "image processor temporal_patch_size", + ) + if ( + processor_patch_size != vision_patch_size + or processor_temporal_patch_size != vision_temporal_patch_size + ): + raise StarVLAError("official Qwen vision and image-processor patch contracts disagree") + contract = { + "channel_count": channel_count, + "patch_size": vision_patch_size, + "temporal_patch_size": vision_temporal_patch_size, + "pixel_patch_width": ( + channel_count * vision_temporal_patch_size * vision_patch_size * vision_patch_size + ), + } + expected = { + "channel_count": 3, + "patch_size": 16, + "temporal_patch_size": 2, + "pixel_patch_width": 1536, + } + if contract != expected: + raise StarVLAError(f"official Qwen processor patch contract changed: {contract}") + return contract + + +def _stack_taps(values: Sequence[Any], *, expected_width: int, label: str) -> Any: + import torch + + if len(values) != EXPECTED_LAYER_COUNT: + raise StarVLAError(f"{label} must contain exactly 36 tensors") + first_shape = tuple(values[0].shape) + if len(first_shape) != 3 or first_shape[0] != 1 or first_shape[2] != expected_width: + raise StarVLAError(f"unexpected {label} tensor shape: {first_shape}") + if any(tuple(value.shape) != first_shape for value in values): + raise StarVLAError(f"{label} tensors do not share one shape") + return torch.stack(list(values), dim=0).squeeze(1) + + +def _canonicalize_qwen_discrete_inputs( + qwen_inputs: Mapping[str, Any], +) -> tuple[np.ndarray, np.ndarray, np.ndarray, dict[str, str]]: + input_ids, input_dtype = _tensor_to_array(qwen_inputs["input_ids"]) + attention_mask, mask_dtype = _tensor_to_array(qwen_inputs["attention_mask"]) + image_grid, grid_dtype = _tensor_to_array(qwen_inputs["image_grid_thw"]) + source_dtypes = { + "input_ids": input_dtype, + "attention_mask": mask_dtype, + "image_grid_thw": grid_dtype, + } + drifted = [name for name, dtype in source_dtypes.items() if dtype != "int64"] + if drifted: + raise StarVLAError( + "official Qwen discrete processor inputs must originate as torch.int64: " + + ", ".join(f"{name}={source_dtypes[name]}" for name in drifted) + ) + if input_ids.ndim != 2 or input_ids.shape[0] != 1 or attention_mask.shape != input_ids.shape: + raise StarVLAError("official Qwen token/mask shape changed") + if not np.all((attention_mask == 0) | (attention_mask == 1)): + raise StarVLAError("official Qwen attention_mask contains values outside {0, 1}") + if not np.all(attention_mask == 1): + raise StarVLAError("single-sample official PI_v3 attention_mask must keep every token") + if image_grid.ndim != 2 or image_grid.shape != (1, 3): + raise StarVLAError(f"official Qwen image_grid_thw shape changed: {image_grid.shape}") + if np.any(image_grid <= 0): + raise StarVLAError( + f"official Qwen image_grid_thw contains non-positive values: {image_grid.tolist()}" + ) + return input_ids[0], attention_mask[0].astype(np.bool_), image_grid, source_dtypes + + +def _validate_qwen_pixel_values( + pixel_values: np.ndarray, + *, + source_dtype: Any, + image_grid: np.ndarray, + pixel_patch_width: int, +) -> None: + expected_patch_count = math.prod(int(value) for value in image_grid.flat) + if ( + pixel_values.dtype != np.float32 + or pixel_values.ndim != 2 + or pixel_values.shape[0] != expected_patch_count + or pixel_values.shape[1] != pixel_patch_width + or source_dtype != "float32" + ): + raise StarVLAError( + "official processor pixel_values are not source-FP32 patches matching " + f"image_grid_thw and width {pixel_patch_width}" + ) + + +def _build_arrays( + captures: Mapping[str, Any], + unnormalized: np.ndarray, + *, + pixel_patch_width: int, +) -> tuple[dict[str, np.ndarray], dict[str, Any]]: + arrays: dict[str, np.ndarray] = {} + records: dict[str, Any] = {} + + def add(name: str, value: Any, *, source_dtype: str | None = None) -> None: + if isinstance(value, np.ndarray): + array = np.ascontiguousarray(value) + inferred_dtype = None + else: + array, inferred_dtype = _tensor_to_array(value) + arrays[name] = array + records[name] = _array_record(array, source_dtype=source_dtype or inferred_dtype) + + qwen_inputs = captures["qwen_inputs"] + required_qwen_inputs = {"input_ids", "attention_mask", "image_grid_thw", "pixel_values"} + missing_qwen_inputs = sorted(required_qwen_inputs - set(qwen_inputs)) + if missing_qwen_inputs: + raise StarVLAError( + f"official Qwen preprocessing did not produce required tensors: {missing_qwen_inputs}" + ) + input_ids, attention_mask, image_grid, source_dtypes = _canonicalize_qwen_discrete_inputs( + qwen_inputs + ) + add("input_ids", input_ids, source_dtype=source_dtypes["input_ids"]) + add("attention_mask", attention_mask, source_dtype=source_dtypes["attention_mask"]) + add("image_grid_thw", image_grid, source_dtype=source_dtypes["image_grid_thw"]) + for key, tensor in sorted(qwen_inputs.items()): + if key not in {"input_ids", "attention_mask", "image_grid_thw"}: + add(f"qwen_input__{key}", tensor) + pixel_values = arrays["qwen_input__pixel_values"] + _validate_qwen_pixel_values( + pixel_values, + source_dtype=records["qwen_input__pixel_values"].get("source_dtype"), + image_grid=image_grid, + pixel_patch_width=pixel_patch_width, + ) + add( + "conditioning_taps", + _stack_taps(captures["conditioning_taps"], expected_width=EXPECTED_QWEN_HIDDEN_DIM, label="conditioning taps"), + ) + add( + "raw_qwen_taps", + _stack_taps(captures["raw_qwen_taps"], expected_width=EXPECTED_QWEN_HIDDEN_DIM, label="raw Qwen taps"), + ) + result_norm, result_norm_dtype = _tensor_to_array(captures["result_norm_diagnostic"]) + if result_norm.shape != (1, input_ids.shape[0], EXPECTED_QWEN_HIDDEN_DIM): + raise StarVLAError( + f"official Qwen result_norm diagnostic shape changed: {result_norm.shape}" + ) + add( + "result_norm_diagnostic", + np.ascontiguousarray(result_norm[0]), + source_dtype=result_norm_dtype, + ) + add( + "projected_hidden_taps", + _stack_taps( + captures["projected_hidden_taps"], + expected_width=EXPECTED_PROJECTED_HIDDEN_DIM, + label="projected hidden taps", + ), + ) + add("initial_noise", captures["initial_noise"]) + add("normalized_actions", captures["normalized_actions"]) + add("unnormalized_actions", np.ascontiguousarray(unnormalized, dtype=np.float32)) + return arrays, records + + +def _runtime_record(torch: Any, transformers: Any, device: str, recorder_probe: Mapping[str, Any]) -> dict[str, Any]: + cuda_device = torch.device(device) + index = cuda_device.index if cuda_device.index is not None else torch.cuda.current_device() + properties = torch.cuda.get_device_properties(index) + return { + "python": platform.python_version(), + "platform": platform.platform(), + "torch": torch.__version__, + "torchvision": _distribution_version("torchvision"), + "transformers": transformers.__version__, + "numpy": np.__version__, + "diffusers": _distribution_version("diffusers"), + "tokenizers": _distribution_version("tokenizers"), + "pillow": _distribution_version("Pillow"), + "omegaconf": _distribution_version("omegaconf"), + "accelerate": _distribution_version("accelerate"), + "safetensors": _distribution_version("safetensors"), + "official_environment_freeze": dict(OFFICIAL_ENVIRONMENT_FREEZE), + "cuda_runtime": torch.version.cuda, + "cudnn": torch.backends.cudnn.version(), + "device": str(cuda_device), + "device_name": properties.name, + "compute_capability": [properties.major, properties.minor], + "qwen3vl_recorder_probe": dict(recorder_probe), + } + + +def _copy_inputs(staging: Path, image_paths: Sequence[Path]) -> list[str]: + inputs_dir = staging / "inputs" + inputs_dir.mkdir() + relative_paths = [] + for index, source in enumerate(image_paths): + suffix = source.suffix.lower() if source.suffix else ".img" + destination = inputs_dir / f"image-{index:02d}{suffix}" + shutil.copyfile(source, destination) + relative_paths.append(destination.relative_to(staging).as_posix()) + return relative_paths + + +def _source_asset_hashes(entry: Mapping[str, Any], *, staged: bool = False) -> dict[str, str]: + overrides = entry.get("staged_overrides", {}) if staged else {} + return { + relative: overrides.get(relative, record)["sha256"] + for relative, record in entry["file_hashes"].items() + } + + +def write_golden( + *, + output_dir: Path, + paths: Mapping[str, Any], + framework: Any, + config: Mapping[str, Any], + recorder_probe: Mapping[str, Any], + image_paths: Sequence[Path], + source_image_records: Sequence[Mapping[str, Any]], + task: str, + unnorm_key: str, + captures: Mapping[str, Any], + unnormalized: np.ndarray, +) -> Path: + import torch + import transformers + + output_dir = output_dir.resolve() + if output_dir.exists(): + raise StarVLAError(f"golden output directory already exists: {output_dir}") + output_dir.parent.mkdir(parents=True, exist_ok=True) + processor_patch_contract = _processor_patch_contract(framework) + arrays, array_records = _build_arrays( + captures, + unnormalized, + pixel_patch_width=processor_patch_contract["pixel_patch_width"], + ) + input_ids = arrays["input_ids"] + attention_mask = arrays["attention_mask"] + if captures["framework_instructions"] != [task]: + raise StarVLAError( + f"official PI_v3 framework instruction changed: {captures['framework_instructions']!r}" + ) + model_instruction = expected_model_instruction(config, task) + rendered_prompt = _render_model_prompt(framework, captures["processed_images"], model_instruction) + token_strings = framework.qwen_vl_interface.processor.tokenizer.convert_ids_to_tokens(input_ids.tolist()) + runtime_contract = expected_runtime_contract() + runtime_contract_sha = _sha256_bytes(_canonical_json(runtime_contract)) + identity = { + "schema_version": GOLDEN_SCHEMA_VERSION, + "variant": SUPPORTED_VARIANT, + "checkpoint_sha256": paths["variant"]["checkpoint"]["sha256"], + "starvla_revision": paths["catalog"]["source_revisions"]["starvla"], + "qwen_revision": paths["qwen"]["revision"], + "runtime_contract_sha256": runtime_contract_sha, + "task": task, + "unnorm_key": unnorm_key, + "seed": SEED, + "images": [record["source_sha256"] for record in source_image_records], + } + golden_id = _sha256_bytes(_canonical_json(identity)) + + with tempfile.TemporaryDirectory(prefix=f".{output_dir.name}.", dir=output_dir.parent) as temporary: + staging = Path(temporary) + copied_images = _copy_inputs(staging, image_paths) + tensor_path = staging / "tensors.npz" + np.savez(tensor_path, **arrays) + image_records = [] + for index, record in enumerate(source_image_records): + copied = staging / copied_images[index] + image_records.append( + { + **record, + "artifact": copied_images[index], + "artifact_size": copied.stat().st_size, + "artifact_sha256": sha256_file(copied), + } + ) + + variant = paths["variant"] + qwen = paths["qwen"] + manifest: dict[str, Any] = { + "schema_version": GOLDEN_SCHEMA_VERSION, + "kind": GOLDEN_KIND, + "golden_id": golden_id, + "created_utc": dt.datetime.now(dt.timezone.utc).isoformat(), + "variant": SUPPORTED_VARIANT, + "model_type": variant["model_type"], + "source": { + "catalog": str(paths["catalog_path"]), + "catalog_sha256": sha256_file(paths["catalog_path"]), + "bundle_uuid": official_bundle_uuid(variant, paths["catalog"]), + "starvla_repo_revision": paths["catalog"]["source_revisions"]["starvla"], + "starvla_checkout": str(paths["source_dir"]), + "checkpoint_repo_id": variant["repo_id"], + "checkpoint_revision": variant["revision"], + "checkpoint_path": str(paths["checkpoint"]), + "checkpoint_size": variant["checkpoint"]["size"], + "checkpoint_sha256": variant["checkpoint"]["sha256"], + "policy_assets": _source_asset_hashes(variant), + "qwen_repo_id": qwen["repo_id"], + "qwen_revision": qwen["revision"], + "qwen_runtime_assets": _source_asset_hashes(qwen), + "qwen_converted_component_assets": _source_asset_hashes(qwen, staged=True), + }, + "runtime": _runtime_record( + torch, transformers, str(next(framework.parameters()).device), recorder_probe + ), + "determinism": { + "seed": SEED, + "rng_reset_immediately_before_predict": True, + "initial_noise_saved_explicitly": True, + "torch_deterministic_algorithms": True, + "cublas_workspace_config": os.environ.get("CUBLAS_WORKSPACE_CONFIG"), + "cuda_matmul_allow_tf32": False, + "cuda_matmul_allow_bf16_reduced_precision_reduction": False, + "cudnn_allow_tf32": False, + "cudnn_benchmark": False, + "attention_implementation": "sdpa", + }, + "compatibility": { + "qwen_bootstrap": ( + "config-only topology construction; all persistent parameters are then populated by " + "strict loading of the pinned official checkpoint" + ), + "effective_config": "config.yaml with pinned checkpoint-derived PI_v3 compatibility fixes", + "projector_autocast": captures["projector_autocast_contract"], + "policy_boundary_casts": { + "projected_hidden": "BF16 output widened exactly to FP32 at DiT cross-attention input", + "initial_noise": "BF16 torch.randn output widened exactly to FP32 at action encoder input", + "reason": ( + "the released source requests CUDA autocast(dtype=float32), which PyTorch 2.6 disables; " + "these two explicit boundary casts realize its declared FP32 policy path" + ), + }, + }, + "input": { + "task": task, + "unnorm_key": unnorm_key, + "state": None, + "images": image_records, + "processed_images": _processed_image_records(captures["processed_images"]), + }, + "model_contract": { + "framework_class": f"{type(framework).__module__}.{type(framework).__name__}", + "action_horizon": EXPECTED_ACTION_HORIZON, + "action_dim": EXPECTED_ACTION_DIM, + "qwen_hidden_dim": EXPECTED_QWEN_HIDDEN_DIM, + "qwen_layer_count": EXPECTED_LAYER_COUNT, + "projected_hidden_dim": EXPECTED_PROJECTED_HIDDEN_DIM, + "hidden_tuple_indices": list(range(1, 37)), + "conditioning_tap_names": list(CONDITIONING_TAP_NAMES), + "raw_tap_names": list(RAW_TAP_NAMES), + "diagnostic_tap_names": [FINAL_NORM_DIAGNOSTIC_NAME], + "tap_layout": "layer_token_hidden", + "conditioning_semantics": CONDITIONING_SEMANTICS, + "result_norm_role": "golden_only_diagnostic_not_conditioning_or_candidate_gate", + "timestep_ids": EXPECTED_TIMESTEP_IDS, + "initial_noise_dtype": "bfloat16", + "policy_compute_dtype": "float32", + "state_input_active": False, + "runtime_contract": runtime_contract, + "runtime_contract_sha256": runtime_contract_sha, + }, + "prompt": { + "framework_instruction": task, + "model_instruction": model_instruction, + "rendered_chat_template": rendered_prompt, + "action_token_mode": "none", + }, + "processor": { + "image_grid_thw": arrays["image_grid_thw"].tolist(), + "pixel_values_shape": list(arrays["qwen_input__pixel_values"].shape), + "patch_contract": processor_patch_contract, + "qwen_input_array_names": sorted( + name for name in arrays if name.startswith("qwen_input__") + ), + "smart_resize_values_are_observed_not_assumed": True, + }, + "tokens": { + "input_ids": input_ids.tolist(), + "attention_mask": attention_mask.tolist(), + "token_strings": token_strings, + }, + "outputs": { + "normalized_actions": arrays["normalized_actions"].tolist(), + "unnormalized_actions": arrays["unnormalized_actions"].tolist(), + }, + "artifacts": { + "tensors": { + "path": tensor_path.name, + "size": tensor_path.stat().st_size, + "sha256": sha256_file(tensor_path), + "encoding": "numpy_npz_stored", + "arrays": array_records, + } + }, + } + manifest["integrity"] = { + "canonicalization": "utf8_json_sort_keys_compact_excluding_integrity", + "manifest_payload_sha256": _sha256_bytes(_canonical_json(manifest)), + } + manifest_path = staging / "golden.json" + manifest_path.write_text( + json.dumps(manifest, ensure_ascii=False, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + Path(temporary).replace(output_dir) + return output_dir / "golden.json" + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Generate an auditable golden from the pinned official StarVLA Qwen3-VL PI_v3 checkpoint." + ) + parser.add_argument("--image", action="append", default=[], type=Path, help="The single 224x224 RGB image") + parser.add_argument("--task", help="Robot task instruction") + parser.add_argument("--unnorm-key", choices=("oxe_bridge", "oxe_rt1")) + parser.add_argument("--output-dir", type=Path) + parser.add_argument("--checkpoint-root", type=Path, default=Path("ckpts/starvla")) + parser.add_argument("--starvla-source", type=Path, default=None) + parser.add_argument("--device", default="cuda:0") + parser.add_argument( + "--preflight-only", + action="store_true", + help="Verify pinned source/assets/checkpoint/runtime/recorder semantics without allocating the full model", + ) + return parser + + +def _load_images(image_paths: Iterable[Path]) -> tuple[list[Any], list[dict[str, Any]]]: + try: + from PIL import Image + except ImportError as exc: + raise StarVLAError("Pillow is required to load oracle images") from exc + images = [] + records = [] + for path in image_paths: + path = path.resolve() + _ensure_regular_file(path, label="input image") + try: + with Image.open(path) as opened: + opened.load() + image = opened.copy() + except (OSError, ValueError) as exc: + raise StarVLAError(f"failed to decode input image {path}: {exc}") from exc + if image.mode != "RGB": + raise StarVLAError(f"official PI_v3 golden input must already be RGB, got mode {image.mode!r}") + if image.size != (224, 224): + raise StarVLAError(f"official PI_v3 golden input must be exactly 224x224, got {image.size}") + images.append(image) + records.append(_image_record(path, image)) + return images, records + + +def main(argv: Sequence[str] | None = None) -> int: + args = build_parser().parse_args(argv) + _require_isolated_python() + if not args.preflight_only: + missing = [ + name + for name, value in ( + ("--image", args.image), + ("--task", args.task), + ("--unnorm-key", args.unnorm_key), + ("--output-dir", args.output_dir), + ) + if not value + ] + if missing: + raise StarVLAError("golden generation requires " + ", ".join(missing)) + if len(args.image) != 1: + raise StarVLAError("the released PI_v3 checkpoint requires exactly one image") + if not args.task.strip(): + raise StarVLAError("--task must not be empty") + checkpoint_root = args.checkpoint_root.resolve() + output_dir = args.output_dir.resolve() + if output_dir == checkpoint_root or checkpoint_root in output_dir.parents: + raise StarVLAError("--output-dir must not be inside the pinned checkpoint source tree") + + checkpoint_root = args.checkpoint_root.resolve() + source_dir = args.starvla_source or checkpoint_root / "source" / "starvla" + paths = validate_official_inputs( + checkpoint_root=checkpoint_root, + source_dir=source_dir, + catalog_path=DEFAULT_CATALOG, + ) + try: + import torch + import transformers + except ImportError as exc: + raise StarVLAError(f"official StarVLA runtime dependency is missing: {exc}") from exc + validate_runtime_versions( + torch_version=torch.__version__, + torchvision_version=_distribution_version("torchvision"), + transformers_version=transformers.__version__, + numpy_version=np.__version__, + diffusers_version=_distribution_version("diffusers"), + tokenizers_version=_distribution_version("tokenizers"), + pillow_version=_distribution_version("Pillow"), + omegaconf_version=_distribution_version("omegaconf"), + accelerate_version=_distribution_version("accelerate"), + safetensors_version=_distribution_version("safetensors"), + ) + _configure_determinism(torch, seed=SEED, device=args.device) + recorder_probe = verify_transformers_qwen3vl_recorder_semantics(torch, transformers) + expected_runtime_contract() + if args.preflight_only: + verify_official_framework_import(paths) + print("Pinned StarVLA PI_v3 oracle preflight passed.") + return 0 + + images, source_image_records = _load_images(args.image) + framework, config = load_official_framework(paths, device=args.device) + captures = run_official_forward(framework, images=images, task=args.task, seed=SEED) + + source_dir = Path(paths["source_dir"]) + sys.path.insert(0, str(source_dir)) + try: + from deployment.model_server import policy_norm_processor + + _assert_module_origin(policy_norm_processor, source_dir) + normalizer = policy_norm_processor.PolicyNormProcessor( + str(paths["checkpoint"]), unnorm_key=args.unnorm_key + ) + normalized = captures["normalized_actions"] + unnormalized = np.asarray(normalizer.unapply_actions(normalized[0]))[None, ...] + if unnormalized.shape != normalized.shape or not np.isfinite(unnormalized).all(): + raise StarVLAError( + f"official action unnormalization returned invalid values/shape: {unnormalized.shape}" + ) + finally: + if sys.path and sys.path[0] == str(source_dir): + del sys.path[0] + + manifest = write_golden( + output_dir=args.output_dir, + paths=paths, + framework=framework, + config=config, + recorder_probe=recorder_probe, + image_paths=args.image, + source_image_records=source_image_records, + task=args.task, + unnorm_key=args.unnorm_key, + captures=captures, + unnormalized=np.ascontiguousarray(unnormalized, dtype=np.float32), + ) + print(f"Wrote official StarVLA PI_v3 golden: {manifest}") + return 0 + + +if __name__ == "__main__": + try: + raise SystemExit(main()) + except StarVLAError as exc: + raise SystemExit(f"error: {exc}") from exc diff --git a/tools/hf2gguf/starvla/pi_v3_golden_constraints.txt b/tools/hf2gguf/starvla/pi_v3_golden_constraints.txt new file mode 100644 index 0000000..d0011fb --- /dev/null +++ b/tools/hf2gguf/starvla/pi_v3_golden_constraints.txt @@ -0,0 +1,14 @@ +# Numeric runtime frozen by the official PI_v3 checkpoint's W&B run. +# Apply this as constraints on top of the pinned StarVLA requirements. +# HF path: wandb/wandb/run-20260426_011111-enstjn5q/files/requirements.txt +# SHA256: de6b505238663ea8a218620e8a4f99cbcfe1e6e09f347ab26f68fe434f3fb00e +torch==2.6.0 +torchvision==0.21.0 +transformers==4.57.0 +numpy==1.26.4 +diffusers==0.37.1 +tokenizers==0.22.2 +pillow==12.1.1 +omegaconf==2.3.0 +accelerate==1.5.2 +safetensors==0.7.0 From 049dc59b21b06ead332e59343f9b9ad83f87c525 Mon Sep 17 00:00:00 2001 From: JJJYmmm <1650675829@qq.com> Date: Mon, 10 Aug 2026 12:42:56 +0800 Subject: [PATCH 05/11] starvla: add Qwen2.5-VL PI policy --- src/models/starvla/pi_policy.cpp | 1081 +++++++++++++ src/models/starvla/pi_policy.h | 93 ++ .../generate_starvla_qwen25_pi_golden.py | 1338 +++++++++++++++++ 3 files changed, 2512 insertions(+) create mode 100644 src/models/starvla/pi_policy.cpp create mode 100644 src/models/starvla/pi_policy.h create mode 100644 tools/hf2gguf/starvla/generate_starvla_qwen25_pi_golden.py diff --git a/src/models/starvla/pi_policy.cpp b/src/models/starvla/pi_policy.cpp new file mode 100644 index 0000000..26fcc46 --- /dev/null +++ b/src/models/starvla/pi_policy.cpp @@ -0,0 +1,1081 @@ +#include "models/starvla/pi_policy.h" + +#include "ggml-backend.h" +#include "ggml.h" +#include "gguf.h" +#include "models/ggml_backend.h" +#include "models/gguf_loader.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace robotcpp::starvla { + +namespace { + +constexpr size_t kGraphSize = 16384; +struct PIBlockWeights { + ggml_tensor * ada_norm_weight = nullptr; + ggml_tensor * ada_norm_bias = nullptr; + ggml_tensor * query_weight = nullptr; + ggml_tensor * query_bias = nullptr; + ggml_tensor * key_weight = nullptr; + ggml_tensor * key_bias = nullptr; + ggml_tensor * value_weight = nullptr; + ggml_tensor * value_bias = nullptr; + ggml_tensor * attention_output_weight = nullptr; + ggml_tensor * attention_output_bias = nullptr; + ggml_tensor * feed_forward_input_weight = nullptr; + ggml_tensor * feed_forward_input_bias = nullptr; + ggml_tensor * feed_forward_output_weight = nullptr; + ggml_tensor * feed_forward_output_bias = nullptr; +}; + +struct PIWeights { + ggml_tensor * timestep_input_weight = nullptr; + ggml_tensor * timestep_input_bias = nullptr; + ggml_tensor * timestep_output_weight = nullptr; + ggml_tensor * timestep_output_bias = nullptr; + std::vector blocks; + ggml_tensor * state_input_weight = nullptr; + ggml_tensor * state_input_bias = nullptr; + ggml_tensor * state_output_weight = nullptr; + ggml_tensor * state_output_bias = nullptr; + ggml_tensor * action_input_weight = nullptr; + ggml_tensor * action_input_bias = nullptr; + ggml_tensor * action_time_mix_weight = nullptr; + ggml_tensor * action_time_mix_bias = nullptr; + ggml_tensor * action_output_weight = nullptr; + ggml_tensor * action_output_bias = nullptr; + ggml_tensor * velocity_input_weight = nullptr; + ggml_tensor * velocity_input_bias = nullptr; + ggml_tensor * velocity_output_weight = nullptr; + ggml_tensor * velocity_output_bias = nullptr; + ggml_tensor * future_tokens = nullptr; + ggml_tensor * action_position = nullptr; +}; + +int require_key(gguf_context * gguf, const char * key, gguf_type type) { + const int index = gguf_find_key(gguf, key); + if (index < 0) { + throw std::runtime_error(std::string("missing required StarVLA PI GGUF metadata: ") + + key); + } + if (gguf_get_kv_type(gguf, index) != type) { + throw std::runtime_error(std::string("invalid StarVLA PI GGUF metadata type: ") + + key); + } + return index; +} + +std::string require_string(gguf_context * gguf, const char * key) { + return gguf_get_val_str(gguf, require_key(gguf, key, GGUF_TYPE_STRING)); +} + +int require_i32(gguf_context * gguf, const char * key) { + return gguf_get_val_i32(gguf, require_key(gguf, key, GGUF_TYPE_INT32)); +} + +float require_f32(gguf_context * gguf, const char * key) { + return gguf_get_val_f32(gguf, require_key(gguf, key, GGUF_TYPE_FLOAT32)); +} + +bool require_bool(gguf_context * gguf, const char * key) { + return gguf_get_val_bool(gguf, require_key(gguf, key, GGUF_TYPE_BOOL)); +} + +int require_array(gguf_context * gguf, const char * key, gguf_type element_type) { + const int index = require_key(gguf, key, GGUF_TYPE_ARRAY); + if (gguf_get_arr_type(gguf, index) != element_type) { + throw std::runtime_error( + std::string("invalid StarVLA PI GGUF array element type: ") + key); + } + return index; +} + +std::vector require_string_array(gguf_context * gguf, + const char * key) { + const int index = require_array(gguf, key, GGUF_TYPE_STRING); + const size_t count = gguf_get_arr_n(gguf, index); + std::vector result; + result.reserve(count); + for (size_t i = 0; i < count; ++i) { + result.emplace_back(gguf_get_arr_str(gguf, index, i)); + } + return result; +} + +std::vector require_i32_array(gguf_context * gguf, const char * key) { + const int index = require_array(gguf, key, GGUF_TYPE_INT32); + const size_t count = gguf_get_arr_n(gguf, index); + const auto * data = + static_cast(gguf_get_arr_data(gguf, index)); + if (data == nullptr && count != 0) { + throw std::runtime_error(std::string("missing StarVLA PI GGUF array data: ") + + key); + } + return count == 0 ? std::vector() + : std::vector(data, data + count); +} + +std::vector require_f32_array(gguf_context * gguf, const char * key) { + const int index = require_array(gguf, key, GGUF_TYPE_FLOAT32); + const size_t count = gguf_get_arr_n(gguf, index); + const auto * data = static_cast(gguf_get_arr_data(gguf, index)); + if (data == nullptr && count != 0) { + throw std::runtime_error(std::string("missing StarVLA PI GGUF array data: ") + + key); + } + return count == 0 ? std::vector() + : std::vector(data, data + count); +} + +std::vector require_bool_array(gguf_context * gguf, const char * key) { + const int index = require_array(gguf, key, GGUF_TYPE_BOOL); + const size_t count = gguf_get_arr_n(gguf, index); + const auto * data = + static_cast(gguf_get_arr_data(gguf, index)); + if (data == nullptr && count != 0) { + throw std::runtime_error(std::string("missing StarVLA PI GGUF array data: ") + + key); + } + std::vector result(count); + for (size_t i = 0; i < count; ++i) { + result[i] = data[i] != 0 ? 1 : 0; + } + return result; +} + +std::string profile_key(int profile_index, const char * suffix) { + return "starvla.normalization.profile." + std::to_string(profile_index) + "." + + suffix; +} + +bool has_shape(const ggml_tensor * tensor, + std::initializer_list expected) { + if (tensor == nullptr || + static_cast(ggml_n_dims(tensor)) != expected.size()) { + return false; + } + size_t dimension = 0; + for (const int64_t value : expected) { + if (tensor->ne[dimension++] != value) { + return false; + } + } + return true; +} + +const char * mode_name(backend_mode mode) { + switch (mode) { + case backend_mode::cpu: + return "cpu"; + case backend_mode::cuda: + return "cuda"; + case backend_mode::metal: + return "metal"; + } + return "unknown"; +} + +std::vector expected_hidden_tuple_indices(int qwen_layer_count, + int block_count) { + std::vector result; + result.reserve(static_cast(block_count)); + const int first = qwen_layer_count + 1 - block_count; + for (int index = first; index <= qwen_layer_count; ++index) { + result.push_back(index); + } + return result; +} + +class PIGGUFLoader final : public gguf_loader { + public: + PIGGUFLoader(PIPolicyConfig & config, PIWeights & weights) + : config_(config), weights_(weights) {} + + protected: + bool parse_metadata(gguf_context * gguf) override { + if (require_string(gguf, "general.architecture") != "starvla-policy" || + require_i32(gguf, "starvla.schema_version") != 1 || + require_string(gguf, "starvla.framework") != "pi") { + throw std::runtime_error("GGUF is not a supported StarVLA PI policy"); + } + config_.backbone_arch = require_string(gguf, "starvla.backbone.arch"); + config_.bundle_uuid = require_string(gguf, "starvla.bundle.uuid"); + config_.text_filename = + require_string(gguf, "starvla.component.text.filename"); + config_.mmproj_filename = + require_string(gguf, "starvla.component.mmproj.filename"); + if (config_.backbone_arch != "qwen2_5_vl" || + config_.bundle_uuid.empty() || config_.text_filename.empty() || + config_.mmproj_filename.empty()) { + throw std::runtime_error("StarVLA PI bundle metadata is incomplete"); + } + + config_.qwen_hidden_dim = + require_i32(gguf, "starvla.qwen.hidden_size"); + config_.qwen_input_embedding_dim = + require_i32(gguf, "starvla.qwen.input_embedding_size"); + config_.qwen_layer_count = + require_i32(gguf, "starvla.qwen.layer_count"); + config_.qwen_vocab_size = require_i32(gguf, "starvla.qwen.vocab_size"); + config_.cot_template = + require_string(gguf, "starvla.prompt.cot_template"); + config_.qwen_hidden_tuple_indices = + require_i32_array(gguf, "starvla.conditioning.hidden_tuple_indices"); + + config_.image_count = require_i32(gguf, "starvla.image.count"); + config_.image_names = + require_string_array(gguf, "starvla.image.names"); + config_.image_framework_inference_pre_resize_width = + require_i32(gguf, "starvla.image.framework_inference_pre_resize_width"); + config_.image_framework_inference_pre_resize_height = + require_i32(gguf, "starvla.image.framework_inference_pre_resize_height"); + config_.image_processor_min_pixels = + require_i32(gguf, "starvla.image.processor_min_pixels"); + config_.image_processor_max_pixels = + require_i32(gguf, "starvla.image.processor_max_pixels"); + config_.image_patch_size = + require_i32(gguf, "starvla.image.patch_size"); + config_.image_spatial_merge_size = + require_i32(gguf, "starvla.image.spatial_merge_size"); + config_.image_min_token_count = + require_i32(gguf, "starvla.image.min_token_count"); + config_.image_max_token_count = + require_i32(gguf, "starvla.image.max_token_count"); + + config_.dit_width = require_i32(gguf, "starvla.pi.dit_width"); + config_.block_count = require_i32(gguf, "starvla.pi.block_count"); + config_.attention_head_count = + require_i32(gguf, "starvla.pi.attention_head_count"); + config_.attention_head_dim = + require_i32(gguf, "starvla.pi.attention_head_dim"); + config_.cross_attention_dim = + require_i32(gguf, "starvla.pi.cross_attention_dim"); + config_.feed_forward_dim = + require_i32(gguf, "starvla.pi.feed_forward_dim"); + config_.mlp_hidden_dim = + require_i32(gguf, "starvla.pi.mlp_hidden_dimension"); + config_.state_dim = require_i32(gguf, "starvla.state.dimension"); + config_.action_dim = require_i32(gguf, "starvla.action.dimension"); + config_.horizon = require_i32(gguf, "starvla.action.horizon"); + config_.state_token_count = + require_i32(gguf, "starvla.pi.state_token_count"); + config_.future_token_count = + require_i32(gguf, "starvla.pi.future_token_count"); + config_.action_position_count = + require_i32(gguf, "starvla.pi.action_position_count"); + config_.timestep_projection_dim = + require_i32(gguf, "starvla.pi.timestep_projection_dim"); + config_.num_inference_timesteps = + require_i32(gguf, "starvla.pi.num_inference_timesteps"); + config_.ada_norm_epsilon = + require_f32(gguf, "starvla.pi.ada_norm_epsilon"); + config_.euler_dt = require_f32(gguf, "starvla.pi.euler_dt"); + config_.timestep_ids = + require_i32_array(gguf, "starvla.pi.timestep_ids"); + + const std::vector expected_indices = + expected_hidden_tuple_indices(config_.qwen_layer_count, + config_.block_count); + const bool dimensions_valid = + config_.qwen_hidden_dim > 0 && + config_.qwen_input_embedding_dim == config_.qwen_hidden_dim && + config_.qwen_layer_count >= config_.block_count && + config_.qwen_vocab_size > 0 && config_.dit_width > 0 && + config_.dit_width % 2 == 0 && config_.block_count > 0 && + config_.attention_head_count > 0 && + config_.attention_head_dim > 0 && + config_.attention_head_count * config_.attention_head_dim == + config_.dit_width && + config_.cross_attention_dim == config_.qwen_hidden_dim && + config_.feed_forward_dim > 0 && + config_.mlp_hidden_dim > 0 && config_.state_dim > 0 && + config_.action_dim > 0 && config_.horizon > 0 && + config_.state_token_count == 1 && + config_.future_token_count > 0 && + config_.action_position_count >= config_.horizon && + config_.timestep_projection_dim >= 4 && + config_.timestep_projection_dim % 2 == 0 && + config_.num_inference_timesteps > 0 && + config_.timestep_ids.size() == + static_cast(config_.num_inference_timesteps) && + std::isfinite(config_.ada_norm_epsilon) && + config_.ada_norm_epsilon > 0.0f && std::isfinite(config_.euler_dt) && + config_.euler_dt > 0.0f && + config_.qwen_hidden_tuple_indices == expected_indices && + config_.image_count > 0 && + config_.image_names.size() == static_cast(config_.image_count) && + config_.image_framework_inference_pre_resize_width > 0 && + config_.image_framework_inference_pre_resize_height > 0 && + config_.image_processor_min_pixels > 0 && + config_.image_processor_max_pixels >= + config_.image_processor_min_pixels && + config_.image_patch_size > 0 && + config_.image_spatial_merge_size > 0 && + config_.image_min_token_count > 0 && + config_.image_max_token_count >= + config_.image_min_token_count && + !config_.cot_template.empty(); + if (!dimensions_valid) { + throw std::runtime_error( + "StarVLA PI dimensions, hidden taps, or sampler schedule are incompatible"); + } + + NormalizationConfig & normalization = config_.normalization; + normalization.clip_actions = + require_bool(gguf, "starvla.normalization.clip_actions"); + normalization.binary_threshold = + require_f32(gguf, "starvla.normalization.binary_threshold"); + normalization.binary_comparison = + require_string(gguf, "starvla.normalization.binary_comparison"); + normalization.continuous_dimensions = + require_i32_array(gguf, "starvla.action.continuous_dimensions"); + normalization.binary_dimensions = + require_i32_array(gguf, "starvla.action.binary_dimensions"); + const int profile_count = + require_i32(gguf, "starvla.normalization.profile_count"); + const std::vector keys = + require_string_array(gguf, "starvla.normalization.profile_keys"); + if (profile_count <= 0 || + keys.size() != static_cast(profile_count)) { + throw std::runtime_error( + "StarVLA PI normalization profile count is inconsistent"); + } + normalization.profiles.clear(); + normalization.profiles.reserve(static_cast(profile_count)); + for (int i = 0; i < profile_count; ++i) { + NormalizationProfile profile; + profile.key = + require_string(gguf, profile_key(i, "key").c_str()); + profile.action_q01 = + require_f32_array(gguf, profile_key(i, "action_q01").c_str()); + profile.action_q99 = + require_f32_array(gguf, profile_key(i, "action_q99").c_str()); + profile.action_mask = + require_bool_array(gguf, + profile_key(i, "action_mask").c_str()); + if (profile.key != keys[static_cast(i)]) { + throw std::runtime_error( + "StarVLA PI normalization profile order is inconsistent"); + } + normalization.profiles.push_back(std::move(profile)); + } + std::string normalization_error; + if (!validate_normalization_config(normalization, config_.action_dim, + normalization_error)) { + throw std::runtime_error(normalization_error); + } + if (!normalization.clip_actions || + normalization.binary_comparison != "ge") { + throw std::runtime_error( + "StarVLA PI normalization must clip actions and use " + "binary comparison 'ge'"); + } + return true; + } + + bool bind_tensors(ggml_context * ctx_data) override { + auto bind = [&](ggml_tensor *& destination, const std::string & name) { + destination = require_tensor(ctx_data, name); + }; + bind(weights_.timestep_input_weight, + "starvla.policy.pi.timestep.input.weight"); + bind(weights_.timestep_input_bias, + "starvla.policy.pi.timestep.input.bias"); + bind(weights_.timestep_output_weight, + "starvla.policy.pi.timestep.output.weight"); + bind(weights_.timestep_output_bias, + "starvla.policy.pi.timestep.output.bias"); + weights_.blocks.clear(); + weights_.blocks.reserve(static_cast(config_.block_count)); + for (int block = 0; block < config_.block_count; ++block) { + const std::string prefix = + "starvla.policy.pi.block." + std::to_string(block) + "."; + PIBlockWeights current; + bind(current.ada_norm_weight, prefix + "ada_norm.weight"); + bind(current.ada_norm_bias, prefix + "ada_norm.bias"); + bind(current.query_weight, prefix + "attention.query.weight"); + bind(current.query_bias, prefix + "attention.query.bias"); + bind(current.key_weight, prefix + "attention.key.weight"); + bind(current.key_bias, prefix + "attention.key.bias"); + bind(current.value_weight, prefix + "attention.value.weight"); + bind(current.value_bias, prefix + "attention.value.bias"); + bind(current.attention_output_weight, + prefix + "attention.output.weight"); + bind(current.attention_output_bias, + prefix + "attention.output.bias"); + bind(current.feed_forward_input_weight, + prefix + "feed_forward.input.weight"); + bind(current.feed_forward_input_bias, + prefix + "feed_forward.input.bias"); + bind(current.feed_forward_output_weight, + prefix + "feed_forward.output.weight"); + bind(current.feed_forward_output_bias, + prefix + "feed_forward.output.bias"); + weights_.blocks.push_back(current); + } + bind(weights_.state_input_weight, + "starvla.policy.pi.state.input.weight"); + bind(weights_.state_input_bias, "starvla.policy.pi.state.input.bias"); + bind(weights_.state_output_weight, + "starvla.policy.pi.state.output.weight"); + bind(weights_.state_output_bias, "starvla.policy.pi.state.output.bias"); + bind(weights_.action_input_weight, + "starvla.policy.pi.action.input.weight"); + bind(weights_.action_input_bias, "starvla.policy.pi.action.input.bias"); + bind(weights_.action_time_mix_weight, + "starvla.policy.pi.action.time_mix.weight"); + bind(weights_.action_time_mix_bias, + "starvla.policy.pi.action.time_mix.bias"); + bind(weights_.action_output_weight, + "starvla.policy.pi.action.output.weight"); + bind(weights_.action_output_bias, + "starvla.policy.pi.action.output.bias"); + bind(weights_.velocity_input_weight, + "starvla.policy.pi.velocity.input.weight"); + bind(weights_.velocity_input_bias, + "starvla.policy.pi.velocity.input.bias"); + bind(weights_.velocity_output_weight, + "starvla.policy.pi.velocity.output.weight"); + bind(weights_.velocity_output_bias, + "starvla.policy.pi.velocity.output.bias"); + bind(weights_.future_tokens, "starvla.policy.pi.future_tokens.weight"); + bind(weights_.action_position, + "starvla.policy.pi.action_position.weight"); + + const int64_t width = config_.dit_width; + if (!has_shape(weights_.timestep_input_weight, + {config_.timestep_projection_dim, width}) || + !has_shape(weights_.timestep_input_bias, {width}) || + !has_shape(weights_.timestep_output_weight, {width, width}) || + !has_shape(weights_.timestep_output_bias, {width}) || + !has_shape(weights_.state_input_weight, + {config_.state_dim, config_.mlp_hidden_dim}) || + !has_shape(weights_.state_input_bias, {config_.mlp_hidden_dim}) || + !has_shape(weights_.state_output_weight, + {config_.mlp_hidden_dim, width}) || + !has_shape(weights_.state_output_bias, {width}) || + !has_shape(weights_.action_input_weight, + {config_.action_dim, width}) || + !has_shape(weights_.action_input_bias, {width}) || + !has_shape(weights_.action_time_mix_weight, + {2 * width, width}) || + !has_shape(weights_.action_time_mix_bias, {width}) || + !has_shape(weights_.action_output_weight, {width, width}) || + !has_shape(weights_.action_output_bias, {width}) || + !has_shape(weights_.velocity_input_weight, + {width, config_.mlp_hidden_dim}) || + !has_shape(weights_.velocity_input_bias, + {config_.mlp_hidden_dim}) || + !has_shape(weights_.velocity_output_weight, + {config_.mlp_hidden_dim, config_.action_dim}) || + !has_shape(weights_.velocity_output_bias, {config_.action_dim}) || + !has_shape(weights_.future_tokens, + {width, config_.future_token_count}) || + !has_shape(weights_.action_position, + {width, config_.action_position_count})) { + throw std::runtime_error( + "StarVLA PI non-transformer tensor has an incompatible ggml shape"); + } + for (const PIBlockWeights & block : weights_.blocks) { + if (!has_shape(block.ada_norm_weight, {width, 2 * width}) || + !has_shape(block.ada_norm_bias, {2 * width}) || + !has_shape(block.query_weight, {width, width}) || + !has_shape(block.query_bias, {width}) || + !has_shape(block.key_weight, + {config_.cross_attention_dim, width}) || + !has_shape(block.key_bias, {width}) || + !has_shape(block.value_weight, + {config_.cross_attention_dim, width}) || + !has_shape(block.value_bias, {width}) || + !has_shape(block.attention_output_weight, {width, width}) || + !has_shape(block.attention_output_bias, {width}) || + !has_shape(block.feed_forward_input_weight, + {width, config_.feed_forward_dim}) || + !has_shape(block.feed_forward_input_bias, + {config_.feed_forward_dim}) || + !has_shape(block.feed_forward_output_weight, + {config_.feed_forward_dim, width}) || + !has_shape(block.feed_forward_output_bias, {width})) { + throw std::runtime_error( + "StarVLA PI transformer tensor has an incompatible ggml shape"); + } + } + return true; + } + + private: + PIPolicyConfig & config_; + PIWeights & weights_; +}; + +std::vector timestep_projection_table(const PIPolicyConfig & config) { + std::vector result( + static_cast(config.num_inference_timesteps) * + config.timestep_projection_dim); + const int half = config.timestep_projection_dim / 2; + for (int step = 0; step < config.num_inference_timesteps; ++step) { + const float timestep = + static_cast(config.timestep_ids[static_cast(step)]); + for (int i = 0; i < half; ++i) { + const float exponent = + -std::log(10000.0f) * i / static_cast(half - 1); + const float angle = timestep * std::exp(exponent); + const size_t offset = + static_cast(step) * config.timestep_projection_dim; + result[offset + static_cast(i)] = std::cos(angle); + result[offset + static_cast(i + half)] = std::sin(angle); + } + } + return result; +} + +std::vector action_time_table(const PIPolicyConfig & config) { + std::vector result( + static_cast(config.num_inference_timesteps) * config.dit_width); + const int half = config.dit_width / 2; + for (int step = 0; step < config.num_inference_timesteps; ++step) { + const float timestep = + static_cast(config.timestep_ids[static_cast(step)]); + for (int i = 0; i < half; ++i) { + const float exponent = + -std::log(10000.0f) * i / static_cast(half); + const float angle = timestep * std::exp(exponent); + const size_t offset = + static_cast(step) * config.dit_width; + result[offset + static_cast(i)] = std::sin(angle); + result[offset + static_cast(i + half)] = std::cos(angle); + } + } + return result; +} + +} // namespace + +struct PIPolicy::Impl { + PIPolicyConfig config; + PIWeights weights; + gguf_load_result loaded; + ggml_backend_t backend_cpu = nullptr; + std::vector backends; + ggml_backend_sched_t scheduler = nullptr; + backend_buft_policy buft_policy; + backend_mode mode = backend_mode::cpu; + int n_threads = 0; + int verbosity = 0; + ggml_context * graph_context = nullptr; + ggml_cgraph * graph = nullptr; + ggml_tensor * hidden_input = nullptr; + ggml_tensor * state_input = nullptr; + ggml_tensor * noise_input = nullptr; + ggml_tensor * timestep_projection_input = nullptr; + ggml_tensor * action_time_input = nullptr; + ggml_tensor * scalar_one_input = nullptr; + ggml_tensor * output = nullptr; + size_t conditioning_token_count = 0; + bool graph_uses_state = false; + size_t graph_builds = 0; + std::vector timestep_table; + std::vector action_table; + + ~Impl() { + clear_graph(); + if (scheduler != nullptr) { + ggml_backend_sched_synchronize(scheduler); + ggml_backend_sched_free(scheduler); + scheduler = nullptr; + } + if (loaded.model_buffer != nullptr) { + ggml_backend_buffer_free(loaded.model_buffer); + loaded.model_buffer = nullptr; + } + if (loaded.ctx_data != nullptr) { + ggml_free(loaded.ctx_data); + loaded.ctx_data = nullptr; + } + if (loaded.gguf != nullptr) { + gguf_free(loaded.gguf); + loaded.gguf = nullptr; + } + for (ggml_backend_t backend : backends) { + if (backend != nullptr) { + ggml_backend_free(backend); + } + } + backends.clear(); + backend_cpu = nullptr; + } + + void clear_graph() { + if (scheduler != nullptr) { + ggml_backend_sched_synchronize(scheduler); + ggml_backend_sched_reset(scheduler); + } + if (graph_context != nullptr) { + ggml_free(graph_context); + graph_context = nullptr; + } + graph = nullptr; + hidden_input = nullptr; + state_input = nullptr; + noise_input = nullptr; + timestep_projection_input = nullptr; + action_time_input = nullptr; + scalar_one_input = nullptr; + output = nullptr; + conditioning_token_count = 0; + graph_uses_state = false; + } + + void build_graph(size_t token_count, bool include_state) { + clear_graph(); + if (token_count == 0 || + token_count > static_cast(std::numeric_limits::max())) { + throw std::runtime_error( + "invalid StarVLA PI conditioning token count"); + } + + ggml_init_params params{}; + params.mem_size = + kGraphSize * ggml_tensor_overhead() + + ggml_graph_overhead_custom(kGraphSize, false); + params.mem_buffer = nullptr; + params.no_alloc = true; + graph_context = ggml_init(params); + if (graph_context == nullptr) { + throw std::runtime_error( + "failed to initialize StarVLA PI graph context"); + } + + const int width = config.dit_width; + const int heads = config.attention_head_count; + const int head_dim = config.attention_head_dim; + hidden_input = ggml_new_tensor_3d( + graph_context, GGML_TYPE_F32, config.qwen_hidden_dim, + static_cast(token_count), config.block_count); + if (include_state) { + state_input = ggml_new_tensor_1d(graph_context, GGML_TYPE_F32, + config.state_dim); + } + noise_input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, + config.action_dim, config.horizon); + timestep_projection_input = ggml_new_tensor_2d( + graph_context, GGML_TYPE_F32, config.timestep_projection_dim, + config.num_inference_timesteps); + action_time_input = ggml_new_tensor_2d( + graph_context, GGML_TYPE_F32, width, + config.num_inference_timesteps); + scalar_one_input = + ggml_new_tensor_1d(graph_context, GGML_TYPE_F32, 1); + if (hidden_input == nullptr || + (include_state && state_input == nullptr) || + noise_input == nullptr || timestep_projection_input == nullptr || + action_time_input == nullptr || scalar_one_input == nullptr) { + throw std::runtime_error( + "failed to create StarVLA PI graph inputs"); + } + ggml_set_name(hidden_input, "starvla_pi_qwen_hidden_states"); + if (state_input != nullptr) { + ggml_set_name(state_input, "starvla_pi_state"); + } + ggml_set_name(noise_input, "starvla_pi_initial_noise"); + ggml_set_name(timestep_projection_input, + "starvla_pi_timestep_projection_table"); + ggml_set_name(action_time_input, "starvla_pi_action_time_table"); + ggml_set_name(scalar_one_input, "starvla_pi_scalar_one"); + ggml_set_input(hidden_input); + if (state_input != nullptr) { + ggml_set_input(state_input); + } + ggml_set_input(noise_input); + ggml_set_input(timestep_projection_input); + ggml_set_input(action_time_input); + ggml_set_input(scalar_one_input); + + auto f32 = [&](ggml_tensor * tensor) { + return tensor->type == GGML_TYPE_F32 + ? tensor + : ggml_cast(graph_context, tensor, GGML_TYPE_F32); + }; + auto linear = [&](ggml_tensor * value, ggml_tensor * weight, + ggml_tensor * bias) { + ggml_tensor * projected = + ggml_mul_mat(graph_context, weight, value); + ggml_mul_mat_set_prec(projected, GGML_PREC_F32); + return ggml_add(graph_context, projected, f32(bias)); + }; + auto ada_norm = [&](ggml_tensor * value, ggml_tensor * temb, + const PIBlockWeights & block) { + ggml_tensor * modulation = + linear(ggml_silu(graph_context, temb), + block.ada_norm_weight, block.ada_norm_bias); + ggml_tensor * scale = + ggml_view_1d(graph_context, modulation, width, 0); + ggml_tensor * shift = ggml_view_1d( + graph_context, modulation, width, + static_cast(width) * sizeof(float)); + ggml_tensor * normalized = + ggml_norm(graph_context, value, config.ada_norm_epsilon); + return ggml_add( + graph_context, + ggml_mul(graph_context, normalized, + ggml_add(graph_context, scale, scalar_one_input)), + shift); + }; + auto attention = [&](ggml_tensor * query_source, + ggml_tensor * key_value_source, + const PIBlockWeights & block) { + const int64_t query_count = query_source->ne[1]; + const int64_t key_value_count = key_value_source->ne[1]; + ggml_tensor * query = + linear(query_source, block.query_weight, block.query_bias); + ggml_tensor * key = + linear(key_value_source, block.key_weight, block.key_bias); + ggml_tensor * value = + linear(key_value_source, block.value_weight, block.value_bias); + query = ggml_reshape_3d(graph_context, query, head_dim, heads, + query_count); + key = ggml_reshape_3d(graph_context, key, head_dim, heads, + key_value_count); + value = ggml_reshape_3d(graph_context, value, head_dim, heads, + key_value_count); + query = ggml_permute(graph_context, query, 0, 2, 1, 3); + key = ggml_permute(graph_context, key, 0, 2, 1, 3); + value = ggml_cont( + graph_context, + ggml_permute(graph_context, value, 1, 2, 0, 3)); + ggml_tensor * scores = ggml_mul_mat(graph_context, key, query); + ggml_mul_mat_set_prec(scores, GGML_PREC_F32); + scores = ggml_soft_max_ext( + graph_context, scores, nullptr, + 1.0f / std::sqrt(static_cast(head_dim)), 0.0f); + ggml_tensor * attended = + ggml_mul_mat(graph_context, value, scores); + ggml_mul_mat_set_prec(attended, GGML_PREC_F32); + attended = + ggml_permute(graph_context, attended, 0, 2, 1, 3); + attended = + ggml_cont_2d(graph_context, attended, width, query_count); + return linear(attended, block.attention_output_weight, + block.attention_output_bias); + }; + + ggml_tensor * state_features = nullptr; + if (include_state) { + state_features = + ggml_relu(graph_context, + linear(state_input, weights.state_input_weight, + weights.state_input_bias)); + state_features = + linear(state_features, weights.state_output_weight, + weights.state_output_bias); + state_features = + ggml_reshape_2d(graph_context, state_features, width, 1); + } + ggml_tensor * future = f32(weights.future_tokens); + ggml_tensor * position_view = ggml_view_2d( + graph_context, weights.action_position, width, config.horizon, + weights.action_position->nb[1], 0); + ggml_tensor * position = f32(position_view); + ggml_tensor * actions = noise_input; + + for (int step = 0; step < config.num_inference_timesteps; ++step) { + ggml_tensor * timestep_projection = ggml_view_1d( + graph_context, timestep_projection_input, + config.timestep_projection_dim, + static_cast(step) * + config.timestep_projection_dim * sizeof(float)); + ggml_tensor * temb = + linear(timestep_projection, weights.timestep_input_weight, + weights.timestep_input_bias); + temb = ggml_silu(graph_context, temb); + temb = linear(temb, weights.timestep_output_weight, + weights.timestep_output_bias); + + ggml_tensor * action_features = + linear(actions, weights.action_input_weight, + weights.action_input_bias); + ggml_tensor * action_time = ggml_view_1d( + graph_context, action_time_input, width, + static_cast(step) * width * sizeof(float)); + action_time = + ggml_repeat(graph_context, action_time, action_features); + action_features = + ggml_concat(graph_context, action_features, action_time, 0); + action_features = + linear(action_features, weights.action_time_mix_weight, + weights.action_time_mix_bias); + action_features = ggml_silu(graph_context, action_features); + action_features = + linear(action_features, weights.action_output_weight, + weights.action_output_bias); + action_features = + ggml_add(graph_context, action_features, position); + + ggml_tensor * hidden = future; + if (state_features != nullptr) { + hidden = + ggml_concat(graph_context, state_features, hidden, 1); + } + hidden = ggml_concat(graph_context, hidden, action_features, 1); + for (int block_index = 0; block_index < config.block_count; + ++block_index) { + const PIBlockWeights & block = + weights.blocks[static_cast(block_index)]; + ggml_tensor * layer_hidden = ggml_view_2d( + graph_context, hidden_input, config.qwen_hidden_dim, + static_cast(token_count), hidden_input->nb[1], + static_cast(block_index) * hidden_input->nb[2]); + ggml_tensor * normalized = ada_norm(hidden, temb, block); + hidden = ggml_add( + graph_context, hidden, + attention(normalized, layer_hidden, block)); + ggml_tensor * ff = + ggml_norm(graph_context, hidden, + config.ada_norm_epsilon); + ff = linear(ff, block.feed_forward_input_weight, + block.feed_forward_input_bias); + ff = ggml_gelu(graph_context, ff); + ff = linear(ff, block.feed_forward_output_weight, + block.feed_forward_output_bias); + hidden = ggml_add(graph_context, hidden, ff); + } + + hidden = + ggml_relu(graph_context, + linear(hidden, weights.velocity_input_weight, + weights.velocity_input_bias)); + hidden = linear(hidden, weights.velocity_output_weight, + weights.velocity_output_bias); + ggml_tensor * velocity = ggml_view_2d( + graph_context, hidden, config.action_dim, config.horizon, + hidden->nb[1], + static_cast( + (include_state ? config.state_token_count : 0) + + config.future_token_count) * + hidden->nb[1]); + actions = ggml_add( + graph_context, actions, + ggml_scale(graph_context, velocity, config.euler_dt)); + } + + output = actions; + ggml_set_name(output, "starvla_pi_normalized_actions"); + ggml_set_output(output); + graph = ggml_new_graph_custom(graph_context, kGraphSize, false); + if (graph == nullptr) { + throw std::runtime_error( + "failed to create StarVLA PI graph"); + } + ggml_build_forward_expand(graph, output); + ggml_backend_sched_reset(scheduler); + if (!ggml_backend_sched_alloc_graph(scheduler, graph)) { + throw std::runtime_error( + "failed to allocate StarVLA PI graph"); + } + conditioning_token_count = token_count; + graph_uses_state = include_state; + ++graph_builds; + } +}; + +PIPolicy::PIPolicy(std::unique_ptr impl) : impl_(std::move(impl)) {} + +PIPolicy::~PIPolicy() = default; + +std::unique_ptr PIPolicy::load(const std::string & path, int n_threads, + int verbosity, std::string & error) { + error.clear(); + if (path.empty()) { + error = "StarVLA PI policy path is required"; + return nullptr; + } + + std::unique_ptr impl(new Impl()); + impl->n_threads = n_threads; + impl->verbosity = verbosity; + try { + backend_scheduler_config scheduler_config; + scheduler_config.max_nodes = static_cast(kGraphSize); + scheduler_config.parallel = false; + scheduler_config.op_offload = true; + backend_loader backend; + if (!backend.load(impl->backend_cpu, impl->backends, impl->scheduler, + impl->buft_policy, true, scheduler_config, verbosity)) { + error = "failed to initialize StarVLA PI backend: " + backend.error(); + return nullptr; + } + impl->mode = backend.mode(); + + PIGGUFLoader loader(impl->config, impl->weights); + if (!loader.load(path.c_str(), impl->buft_policy.model_buft, + impl->loaded, verbosity)) { + error = loader.error(); + return nullptr; + } + if (impl->loaded.ctx_data == nullptr || + impl->loaded.model_buffer == nullptr) { + error = "StarVLA PI policy GGUF has no tensors"; + return nullptr; + } + ggml_backend_buffer_set_usage( + impl->loaded.model_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + impl->timestep_table = timestep_projection_table(impl->config); + impl->action_table = action_time_table(impl->config); + if (verbosity >= 1) { + std::fprintf( + stderr, + "%s: backend=%s qwen=%d width=%d blocks=%d horizon=%d " + "action_dim=%d profiles=%zu\n", + __func__, mode_name(impl->mode), impl->config.qwen_hidden_dim, + impl->config.dit_width, impl->config.block_count, + impl->config.horizon, impl->config.action_dim, + impl->config.normalization.profiles.size()); + } + } catch (const std::exception & exception) { + error = exception.what(); + return nullptr; + } + return std::unique_ptr(new PIPolicy(std::move(impl))); +} + +bool PIPolicy::evaluate(const float * qwen_hidden_states, + size_t hidden_element_count, const float * state, + size_t state_element_count, const float * initial_noise, + size_t noise_element_count, + std::vector & normalized_actions, + std::string & error) { + normalized_actions.clear(); + error.clear(); + if (impl_ == nullptr || impl_->scheduler == nullptr) { + error = "StarVLA PI policy is not initialized"; + return false; + } + const size_t layer_width = + static_cast(impl_->config.block_count) * + impl_->config.qwen_hidden_dim; + if (qwen_hidden_states == nullptr || layer_width == 0 || + hidden_element_count == 0 || + hidden_element_count % layer_width != 0) { + error = + "StarVLA PI layer-wise Qwen conditioning tensor has an incompatible shape"; + return false; + } + const size_t token_count = hidden_element_count / layer_width; + if (token_count == 0 || + token_count > static_cast(std::numeric_limits::max())) { + error = + "StarVLA PI layer-wise Qwen conditioning tensor has an incompatible shape"; + return false; + } + const bool include_state = state_element_count != 0; + if (include_state && + (state == nullptr || + state_element_count != static_cast(impl_->config.state_dim))) { + error = "StarVLA PI state tensor has an incompatible shape"; + return false; + } + const size_t expected_noise = + static_cast(impl_->config.horizon) * impl_->config.action_dim; + if (initial_noise == nullptr || noise_element_count != expected_noise) { + error = "StarVLA PI initial-noise tensor has an incompatible shape"; + return false; + } + if (std::any_of(qwen_hidden_states, + qwen_hidden_states + hidden_element_count, + [](float value) { return !std::isfinite(value); }) || + (include_state && + std::any_of(state, state + state_element_count, + [](float value) { return !std::isfinite(value); })) || + std::any_of(initial_noise, initial_noise + noise_element_count, + [](float value) { return !std::isfinite(value); })) { + error = + "StarVLA PI conditioning, state, and initial noise must be finite"; + return false; + } + + try { + if (impl_->graph == nullptr || + impl_->conditioning_token_count != token_count || + impl_->graph_uses_state != include_state) { + impl_->build_graph(token_count, include_state); + } + } catch (const std::exception & exception) { + error = exception.what(); + return false; + } + + ggml_backend_tensor_set(impl_->hidden_input, qwen_hidden_states, 0, + hidden_element_count * sizeof(float)); + if (include_state) { + ggml_backend_tensor_set(impl_->state_input, state, 0, + state_element_count * sizeof(float)); + } + ggml_backend_tensor_set(impl_->noise_input, initial_noise, 0, + noise_element_count * sizeof(float)); + ggml_backend_tensor_set(impl_->timestep_projection_input, + impl_->timestep_table.data(), 0, + impl_->timestep_table.size() * sizeof(float)); + ggml_backend_tensor_set(impl_->action_time_input, + impl_->action_table.data(), 0, + impl_->action_table.size() * sizeof(float)); + const float one = 1.0f; + ggml_backend_tensor_set(impl_->scalar_one_input, &one, 0, sizeof(one)); + set_backend_threads(impl_->backends, impl_->n_threads); + if (ggml_backend_sched_graph_compute(impl_->scheduler, impl_->graph) != + GGML_STATUS_SUCCESS) { + error = "StarVLA PI graph compute failed"; + return false; + } + + normalized_actions.resize(expected_noise); + ggml_backend_tensor_get(impl_->output, normalized_actions.data(), 0, + expected_noise * sizeof(float)); + if (std::any_of(normalized_actions.begin(), normalized_actions.end(), + [](float value) { return !std::isfinite(value); })) { + normalized_actions.clear(); + error = "StarVLA PI graph produced non-finite actions"; + return false; + } + return true; +} + +bool PIPolicy::unnormalize(const std::vector & normalized_actions, + const std::string & profile_key_value, + std::vector & actions, + std::string & error) const { + if (impl_ == nullptr) { + actions.clear(); + error = "StarVLA PI policy is not initialized"; + return false; + } + return denormalize_actions(impl_->config.normalization, profile_key_value, + normalized_actions, impl_->config.horizon, + impl_->config.action_dim, actions, error); +} + +const PIPolicyConfig & PIPolicy::config() const { + if (impl_ == nullptr) { + throw std::runtime_error("StarVLA PI policy is not initialized"); + } + return impl_->config; +} + +const char * PIPolicy::backend_name() const { + return impl_ != nullptr ? mode_name(impl_->mode) : "unknown"; +} + +size_t PIPolicy::graph_build_count() const { + return impl_ != nullptr ? impl_->graph_builds : 0; +} + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/pi_policy.h b/src/models/starvla/pi_policy.h new file mode 100644 index 0000000..6f6923f --- /dev/null +++ b/src/models/starvla/pi_policy.h @@ -0,0 +1,93 @@ +#pragma once + +#include "models/starvla/normalization.h" + +#include +#include +#include +#include +#include + +namespace robotcpp::starvla { + +struct PIPolicyConfig { + std::string backbone_arch; + std::string bundle_uuid; + std::string text_filename; + std::string mmproj_filename; + + int qwen_hidden_dim = 0; + int qwen_input_embedding_dim = 0; + int qwen_layer_count = 0; + int qwen_vocab_size = 0; + std::string cot_template; + std::vector qwen_hidden_tuple_indices; + int image_count = 0; + std::vector image_names; + int image_framework_inference_pre_resize_width = 0; + int image_framework_inference_pre_resize_height = 0; + int image_processor_min_pixels = 0; + int image_processor_max_pixels = 0; + int image_patch_size = 0; + int image_spatial_merge_size = 0; + int image_min_token_count = 0; + int image_max_token_count = 0; + + int dit_width = 0; + int block_count = 0; + int attention_head_count = 0; + int attention_head_dim = 0; + int cross_attention_dim = 0; + int feed_forward_dim = 0; + int mlp_hidden_dim = 0; + int state_dim = 0; + int action_dim = 0; + int horizon = 0; + int state_token_count = 0; + int future_token_count = 0; + int action_position_count = 0; + int timestep_projection_dim = 0; + int num_inference_timesteps = 0; + float ada_norm_epsilon = 0.0f; + float euler_dt = 0.0f; + std::vector timestep_ids; + NormalizationConfig normalization; +}; + +class PIPolicy { + public: + ~PIPolicy(); + + PIPolicy(const PIPolicy &) = delete; + PIPolicy & operator=(const PIPolicy &) = delete; + + static std::unique_ptr load(const std::string & path, int n_threads, + int verbosity, std::string & error); + + // qwen_hidden_states is layer-major + // [block_count, token_count, qwen_hidden_dim]. The legacy released + // implementation did not forward the Qwen attention mask into the policy + // head. state is either omitted (the official Bridge deployment path) or + // one token [state_dim], and initial_noise is token-major + // [horizon, action_dim]. + bool evaluate(const float * qwen_hidden_states, size_t hidden_element_count, + const float * state, size_t state_element_count, + const float * initial_noise, size_t noise_element_count, + std::vector & normalized_actions, std::string & error); + bool unnormalize(const std::vector & normalized_actions, + const std::string & profile_key, std::vector & actions, + std::string & error) const; + + const PIPolicyConfig & config() const; + const char * backend_name() const; + size_t graph_build_count() const; + + private: + struct Impl; + + explicit PIPolicy(std::unique_ptr impl); + + std::unique_ptr impl_; +}; + +} // namespace robotcpp::starvla diff --git a/tools/hf2gguf/starvla/generate_starvla_qwen25_pi_golden.py b/tools/hf2gguf/starvla/generate_starvla_qwen25_pi_golden.py new file mode 100644 index 0000000..b83e915 --- /dev/null +++ b/tools/hf2gguf/starvla/generate_starvla_qwen25_pi_golden.py @@ -0,0 +1,1338 @@ +#!/usr/bin/env python3 +"""Generate a fixed-noise local-Python oracle for the released Qwen2.5 PI. + +The published checkpoint predates the current QwenPI refactor. This exporter +therefore executes the exact historical implementation stored in the pinned +local StarVLA git repository, including the documented ``--use_bf16`` +deployment path. It applies one bootstrap shim: + +* construct Qwen2.5-VL from its local config before the complete checkpoint is + loaded, avoiding a duplicate base-weight download. + +After strict loading, the whole framework is converted to BF16 exactly as in +the official server command. The action head remains the historical +16-block, all-cross-attention forward. +Its initial 16x7 noise tensor is an explicit binary-fraction fixture shared +with the C++ parity runner; cross-language RNG replay is never used. +""" + +from __future__ import annotations + +import argparse +import contextlib +import datetime as dt +import gc +import hashlib +import io +import json +import math +import os +import shutil +import subprocess +import sys +import tarfile +import tempfile +from pathlib import Path +from typing import Any, Iterable, Mapping, Sequence + +import numpy as np + + +TOOLS_DIR = Path(__file__).resolve().parent +if str(TOOLS_DIR) not in sys.path: + sys.path.insert(0, str(TOOLS_DIR)) + +from generate_starvla_oft_golden import ( # noqa: E402 + _assert_module_origin, + _canonical_json, + _configure_determinism, + _distribution_version, + _ensure_regular_file, + _image_pixel_sha256, + _require_isolated_python, + _runtime_record, + _sha256_bytes, + validate_runtime_versions, +) +from generate_starvla_qwen25_groot_golden import ( # noqa: E402 + ACTION_TOKEN_COUNT, + ACTION_TOKEN_ID_MAX, + ACTION_TOKEN_ID_MIN, + EXPECTED_QWEN_VL_UTILS_VERSION, + validate_action_tokenizer_assets, + validate_processor_contract, +) +from generate_starvla_qwen25_oft_golden import ( # noqa: E402 + _official_qwen25_alias, + _verify_clean_source, +) +from starvla_checkpoint import ( # noqa: E402 + DEFAULT_CATALOG, + StarVLAError, + get_variant, + load_catalog, + official_bundle_uuid, + sha256_file, + verify_catalog_files, + verify_checkpoint_file, +) + + +SCHEMA_VERSION = 1 +GOLDEN_KIND = "starvla_qwen25_pi_local_pt_python_oracle" +MODEL_TYPE = "starvla" +VARIANT = "qwen25_pi" +BACKBONE = "qwen2_5_vl" +ACTION_RELATIVE_L2_LIMIT = 0.03 + +OFFICIAL_CHECKPOINT_REPO_ID = "StarVLA/Qwen-PI-Bridge-RT-1" +OFFICIAL_CHECKPOINT_REVISION = "26d0e079fbe3bc3fc62301f44f0025ef7c64ee22" +OFFICIAL_CHECKPOINT_FILENAME = "steps_30000_pytorch_model.pt" +OFFICIAL_CHECKPOINT_SIZE = 10_103_104_403 +OFFICIAL_CHECKPOINT_SHA256 = ( + "8a0e47858921924d5038f7c4393dee6682b83175a85546e35e357e8f74ce8343" +) +OFFICIAL_QWEN_REPO_ID = "StarVLA/Qwen2.5-VL-3B-Instruct-Action" +OFFICIAL_QWEN_REVISION = "ce86bd9a53416527b8361e8dfc47316288ffa110" +OFFICIAL_STARVLA_REPO_ID = "starVLA/starVLA" +OFFICIAL_STARVLA_REVISION = "631aae02afe6d95876e923ff518e8ff2ab9a2f88" +LEGACY_IMPLEMENTATION_REVISION = "e872a8579055f9332add8a2549b9fd5599e11510" +PI_RUNTIME_CONTRACT_SHA256 = ( + "dea02dbb9099b34454db473c39375ce6467109287a27d4ab89193561be035219" +) + +EXPECTED_ACTION_HORIZON = 16 +EXPECTED_ACTION_DIM = 7 +EXPECTED_STATE_DIM = 7 +EXPECTED_QWEN_HIDDEN_DIM = 2048 +EXPECTED_QWEN_LAYER_COUNT = 36 +EXPECTED_HIDDEN_TUPLE_INDICES = list(range(21, 37)) +EXPECTED_DIT_BLOCK_COUNT = 16 +EXPECTED_DIT_WIDTH = 2048 +EXPECTED_FUTURE_TOKEN_COUNT = 32 +EXPECTED_TIMESTEP_IDS = [0, 250, 500, 750] +EXPECTED_COT_TEMPLATE = ( + "Your task is {instruction}. To identify the key objects for your task. " + "Locate their bounding boxes in [x1,y1,x2,y2] format." +) +UNNORM_KEYS = ("oxe_bridge", "oxe_rt1") + +NOISE_ALGORITHM = "portable_binary_fraction_lcg_v1" +NOISE_DENOMINATOR = 64 +NOISE_MULTIPLIER = 73 +NOISE_INCREMENT = 19 +NOISE_MODULUS = 257 +NOISE_OFFSET = 128 + +LEGACY_SOURCE_FILES = { + "deployment/model_server/README.md": + "85662206d8f9ba1948ccc2c588b241fe9f45e0ee43e77ce825a2247f195cb3a6", + "deployment/model_server/server_policy.py": + "98569a4d3a1781d9c9b0fa5bd1952c4f212ff5307522078544bee59059a7df17", + "examples/LIBERO/model2libero_interface.py": + "16f760d011513f6be4f6fbc304aa6567f0a517015122f8c7f67ac85a07f37a13", + "examples/SimplerEnv/model2simpler_interface.py": + "510ae919871ccd6ce64271338c9dfb01648a3c0a820adf90998537ed9bf0fac3", + "starVLA/model/framework/base_framework.py": + "12cdfc8afbff72a44e3f4d0bbafc229721e49de26fe97c5b79821db06118c334", + "starVLA/model/framework/QwenPI.py": + "d368c669ec178045ca4143c7f90c6db75082946042afe915d4686e18d43be525", + "starVLA/model/modules/action_model/LayerwiseFM_ActionHeader.py": + "c586021a5d98605c01728d3ccc98218ce3bc639a95c06f1579eb8289612f5d43", + "starVLA/model/modules/action_model/flow_matching_head/cross_attention_dit.py": + "d835796a351f4562b826ada959332c7baa063b79432ca15e4d0ce76745128a62", + "starVLA/model/modules/vlm/QWen2_5.py": + "b94b9220a04ad6017789e9e907fc2d6e7ec8c32f77a7a38d4adebd18bf2fe5c3", +} + + +def explicit_initial_noise() -> np.ndarray: + """Return the portable 16x7 parity noise using exact binary fractions.""" + + count = EXPECTED_ACTION_HORIZON * EXPECTED_ACTION_DIM + index = np.arange(count, dtype=np.int64) + numerator = ( + (index * NOISE_MULTIPLIER + NOISE_INCREMENT) % NOISE_MODULUS + ) - NOISE_OFFSET + result = numerator.astype(np.float32) / np.float32(NOISE_DENOMINATOR) + return np.ascontiguousarray( + result.reshape(1, EXPECTED_ACTION_HORIZON, EXPECTED_ACTION_DIM) + ) + + +def _load_json_object(path: Path, *, label: str) -> dict[str, Any]: + _ensure_regular_file(path, label=label) + try: + value = json.loads(path.read_text(encoding="utf-8")) + except (OSError, UnicodeError, json.JSONDecodeError) as exc: + raise StarVLAError(f"failed to parse {label} {path}: {exc}") from exc + if not isinstance(value, dict): + raise StarVLAError(f"{label} root must be an object") + return value + + +def _run_git(source_dir: Path, *args: str, binary: bool = False) -> bytes | str: + try: + result = subprocess.run( + ["git", "-C", str(source_dir), *args], + check=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + except (OSError, subprocess.CalledProcessError) as exc: + detail = "" + if isinstance(exc, subprocess.CalledProcessError): + detail = exc.stderr.decode("utf-8", errors="replace").strip() + raise StarVLAError( + f"failed to inspect pinned StarVLA git object: {detail or exc}" + ) from exc + return result.stdout if binary else result.stdout.decode("utf-8").strip() + + +def verify_source_semantics(source_dir: Path) -> dict[str, Any]: + """Bind the exact legacy implementation used by the released checkpoint.""" + + object_type = _run_git( + source_dir, "cat-file", "-t", LEGACY_IMPLEMENTATION_REVISION + ) + if object_type != "commit": + raise StarVLAError( + f"legacy PI revision is not a commit: {LEGACY_IMPLEMENTATION_REVISION}" + ) + + files: dict[str, str] = {} + sources: dict[str, str] = {} + for relative, expected in LEGACY_SOURCE_FILES.items(): + payload = _run_git( + source_dir, + "show", + f"{LEGACY_IMPLEMENTATION_REVISION}:{relative}", + binary=True, + ) + assert isinstance(payload, bytes) + digest = hashlib.sha256(payload).hexdigest() + if digest != expected: + raise StarVLAError( + f"legacy PI source SHA256 mismatch for {relative}: " + f"expected {expected}, got {digest}" + ) + files[relative] = digest + sources[relative] = payload.decode("utf-8") + + framework_source = sources["starVLA/model/framework/QwenPI.py"] + action_source = sources[ + "starVLA/model/modules/action_model/LayerwiseFM_ActionHeader.py" + ] + base_source = sources["starVLA/model/framework/base_framework.py"] + server_source = sources["deployment/model_server/server_policy.py"] + deployment_readme = sources["deployment/model_server/README.md"] + evaluator_sources = ( + sources["examples/LIBERO/model2libero_interface.py"], + sources["examples/SimplerEnv/model2simpler_interface.py"], + ) + required_framework = ( + "expected_layers = len(self.action_model.model.transformer_blocks)", + "vl_embs_list = list(all_hidden[-expected_layers:])", + 'getattr(self.config.datasets.vla_data, "image_size", None)', + 'with torch.autocast("cuda", dtype=torch.float32):', + ) + required_action = ( + "for layer_idx, layer in enumerate(self.model.transformer_blocks):", + "encoder_hidden_states=vl_embs_list[layer_idx]", + "actions = torch.randn(", + "actions = actions + dt * pred_velocity", + ) + required_normalization = ( + "normalized_actions = np.clip(normalized_actions, -1, 1)", + "normalized_actions[:, 6] = np.where(normalized_actions[:, 6] < 0.5, 0, 1)", + ) + missing = [ + fragment + for fragment in required_framework + if fragment not in framework_source + ] + [ + fragment for fragment in required_action if fragment not in action_source + ] + [ + fragment for fragment in required_normalization if fragment not in base_source + ] + [ + fragment + for fragment in ("vla = vla.to(torch.bfloat16)",) + if fragment not in server_source + ] + [ + fragment + for fragment in ("--use_bf16",) + if fragment not in deployment_readme + ] + [ + fragment + for evaluator_source in evaluator_sources + for fragment in required_normalization + if fragment not in evaluator_source + ] + if missing: + raise StarVLAError( + f"legacy PI source semantics probe failed: {missing!r}" + ) + return { + "revision": LEGACY_IMPLEMENTATION_REVISION, + "files": files, + "checkpoint_block_count": EXPECTED_DIT_BLOCK_COUNT, + "hidden_tuple_indices": EXPECTED_HIDDEN_TUPLE_INDICES, + "block_mode": "layerwise_cross_attention_every_block", + "released_config_interleave_self_attention": True, + "use_canonical_dit_forward": False, + "attention_mask_runtime_active": False, + "deployment_precision": "whole_model_bf16_via_use_bf16", + "normalization": "clip_minus1_plus1_then_binary_ge_0_5", + } + + +def _validate_catalog_identity( + catalog: Mapping[str, Any], +) -> tuple[dict[str, Any], dict[str, Any]]: + variant = get_variant(catalog, VARIANT) + qwen_key = variant.get("qwen_asset") + qwen = catalog.get("shared_assets", {}).get(qwen_key) + if not isinstance(qwen, dict): + raise StarVLAError(f"catalog variant {VARIANT} has no Qwen action asset") + expected_checkpoint = { + "path": f"checkpoints/{OFFICIAL_CHECKPOINT_FILENAME}", + "size": OFFICIAL_CHECKPOINT_SIZE, + "sha256": OFFICIAL_CHECKPOINT_SHA256, + } + if ( + variant.get("repo_id") != OFFICIAL_CHECKPOINT_REPO_ID + or variant.get("revision") != OFFICIAL_CHECKPOINT_REVISION + or variant.get("checkpoint") != expected_checkpoint + ): + raise StarVLAError("catalog Qwen2.5 PI checkpoint identity drifted") + if ( + qwen.get("repo_id") != OFFICIAL_QWEN_REPO_ID + or qwen.get("revision") != OFFICIAL_QWEN_REVISION + ): + raise StarVLAError("catalog Qwen2.5 PI action-tokenizer identity drifted") + if ( + catalog.get("source_revisions", {}).get("starvla") + != OFFICIAL_STARVLA_REVISION + ): + raise StarVLAError("catalog StarVLA source revision drifted") + return variant, qwen + + +def _validate_effective_config(config: Mapping[str, Any]) -> None: + try: + framework = config["framework"] + action = framework["action_model"] + diffusion = action["diffusion_model_cfg"] + vla = config["datasets"]["vla_data"] + except (KeyError, TypeError) as exc: + raise StarVLAError("effective Qwen2.5 PI config is incomplete") from exc + actual = { + "framework": framework.get("name"), + "base_vlm": framework.get("qwenvl", {}).get("base_vlm"), + "action_model_type": action.get("action_model_type"), + "configured_hidden_size": action.get("hidden_size"), + "historical_runtime_hidden_size": action.get("action_hidden_dim"), + "action_horizon": action.get("action_horizon"), + "future_action_window_size": action.get("future_action_window_size"), + "action_dim": action.get("action_dim"), + "state_dim": action.get("state_dim"), + "steps": action.get("num_inference_timesteps"), + "buckets": action.get("num_timestep_buckets"), + "future_tokens": action.get("num_target_vision_tokens"), + "layers": diffusion.get("num_layers"), + "cross_dim": diffusion.get("cross_attention_dim"), + "output_dim": diffusion.get("output_dim"), + "interleave": diffusion.get("interleave_self_attention"), + "image_size": vla.get("image_size"), + "default_image_resolution": vla.get("default_image_resolution"), + "obs": vla.get("obs"), + "data_mix": vla.get("data_mix"), + "cot": vla.get("CoT_prompt"), + } + expected = { + "framework": "QwenPI", + "base_vlm": "starVLA/Qwen2.5-VL-3B-Instruct-Action", + "action_model_type": "DiT-Qwen", + "configured_hidden_size": 1024, + "historical_runtime_hidden_size": EXPECTED_DIT_WIDTH, + "action_horizon": EXPECTED_ACTION_HORIZON, + "future_action_window_size": EXPECTED_ACTION_HORIZON - 1, + "action_dim": EXPECTED_ACTION_DIM, + "state_dim": EXPECTED_STATE_DIM, + "steps": 4, + "buckets": 1000, + "future_tokens": EXPECTED_FUTURE_TOKEN_COUNT, + "layers": EXPECTED_DIT_BLOCK_COUNT, + "cross_dim": EXPECTED_QWEN_HIDDEN_DIM, + "output_dim": 1024, + "interleave": True, + "image_size": [224, 224], + "default_image_resolution": [3, 224, 224], + "obs": ["image_0"], + "data_mix": "bridge_rt_1", + "cot": EXPECTED_COT_TEMPLATE, + } + if actual != expected: + raise StarVLAError(f"unexpected effective Qwen2.5 PI config: {actual}") + + +def normalization_contract( + norm_stats: Mapping[str, Any], unnorm_key: str +) -> dict[str, Any]: + if unnorm_key not in UNNORM_KEYS: + raise StarVLAError( + f"Qwen2.5 PI unnorm_key must be one of {list(UNNORM_KEYS)}" + ) + profile = norm_stats.get(unnorm_key) + action = profile.get("action") if isinstance(profile, Mapping) else None + state = profile.get("state") if isinstance(profile, Mapping) else None + if not isinstance(action, Mapping) or not isinstance(state, Mapping): + raise StarVLAError( + f"dataset statistics has no complete {unnorm_key} action/state objects" + ) + try: + q01 = np.asarray(action["q01"], dtype=np.float32) + q99 = np.asarray(action["q99"], dtype=np.float32) + mask = np.asarray(action["mask"], dtype=np.bool_) + state_q01 = np.asarray(state["q01"], dtype=np.float32) + except (KeyError, TypeError, ValueError) as exc: + raise StarVLAError( + f"invalid Qwen2.5 PI statistics for {unnorm_key}: {exc}" + ) from exc + if q01.shape != (7,) or q99.shape != (7,) or mask.shape != (7,): + raise StarVLAError("Qwen2.5 PI action statistics must be 7D") + if state_q01.shape != (8,): + raise StarVLAError( + "Qwen2.5 PI dataset state statistics must remain 8D" + ) + if ( + not np.isfinite(q01).all() + or not np.isfinite(q99).all() + or np.any(q99[mask] <= q01[mask]) + or mask.tolist() != [True, True, True, True, True, True, False] + ): + raise StarVLAError("Qwen2.5 PI normalization statistics are invalid") + return { + "stats_key": unnorm_key, + "q01": q01.tolist(), + "q99": q99.tolist(), + "mask": mask.tolist(), + "continuous_dimensions": [0, 1, 2, 3, 4, 5], + "binary_dimensions": [6], + "binary_threshold": 0.5, + "binary_comparison": "ge", + "continuous_clip": True, + "state_input_contract": + "caller_supplies_model_7d_state_8d_dataset_stats_are_not_applied", + } + + +def unnormalize_actions( + normalized: np.ndarray, + norm_stats: Mapping[str, Any], + unnorm_key: str, +) -> np.ndarray: + contract = normalization_contract(norm_stats, unnorm_key) + values = np.clip( + np.ascontiguousarray(normalized, dtype=np.float32), + np.float32(-1.0), + np.float32(1.0), + ) + if values.shape != (1, EXPECTED_ACTION_HORIZON, EXPECTED_ACTION_DIM): + raise StarVLAError( + f"normalized Qwen2.5 PI actions have invalid shape: {values.shape}" + ) + q01 = np.asarray(contract["q01"], dtype=np.float32) + q99 = np.asarray(contract["q99"], dtype=np.float32) + mask = np.asarray(contract["mask"], dtype=np.bool_) + output = np.empty_like(values) + output[..., mask] = ( + (values[..., mask] + np.float32(1.0)) + * np.float32(0.5) + * (q99[mask] - q01[mask]) + + q01[mask] + ) + output[..., ~mask] = ( + values[..., ~mask] >= np.float32(contract["binary_threshold"]) + ).astype(np.float32) + if not np.isfinite(output).all(): + raise StarVLAError("Qwen2.5 PI unnormalization produced non-finite values") + return np.ascontiguousarray(output) + + +def validate_local_inputs( + *, + checkpoint_root: Path, + checkpoint: Path | None, + qwen_model: Path | None, + source_dir: Path, + catalog_path: Path = DEFAULT_CATALOG, +) -> dict[str, Any]: + catalog = load_catalog(catalog_path) + variant, qwen = _validate_catalog_identity(catalog) + checkpoint_root = checkpoint_root.resolve() + policy_dir = checkpoint_root / "sources" / variant["directory"] / variant["revision"] + qwen_dir = checkpoint_root / "sources" / qwen["directory"] / qwen["revision"] + checkpoint_path = policy_dir / variant["checkpoint"]["path"] + if checkpoint is not None and checkpoint.resolve() != checkpoint_path.resolve(): + raise StarVLAError( + f"Qwen2.5 PI checkpoint must be the catalog path {checkpoint_path}" + ) + if qwen_model is not None and qwen_model.resolve() != qwen_dir.resolve(): + raise StarVLAError( + f"Qwen2.5 PI processor must be the catalog path {qwen_dir}" + ) + + verify_catalog_files(policy_dir, variant) + verify_catalog_files(qwen_dir, qwen) + tokenizer = validate_action_tokenizer_assets(qwen_dir) + source_dir = source_dir.resolve() + revision = _verify_clean_source(source_dir, OFFICIAL_STARVLA_REVISION) + source_probe = verify_source_semantics(source_dir) + + sidecar = Path(f"{checkpoint_path}.aria2") + checkpoint_ready = ( + checkpoint_path.is_file() + and not checkpoint_path.is_symlink() + and not sidecar.exists() + ) + if checkpoint_ready: + verify_checkpoint_file(checkpoint_path, variant) + + config_yaml = policy_dir / "config.yaml" + dataset_statistics = policy_dir / "dataset_statistics.json" + try: + import yaml + + config = yaml.safe_load(config_yaml.read_text(encoding="utf-8")) + except (ImportError, OSError, UnicodeError, ValueError) as exc: + raise StarVLAError(f"failed to load Qwen2.5 PI config.yaml: {exc}") from exc + if not isinstance(config, dict): + raise StarVLAError("Qwen2.5 PI config.yaml root must be an object") + _validate_effective_config(config) + norm_stats = _load_json_object( + dataset_statistics, label="Qwen2.5 PI dataset statistics" + ) + if set(norm_stats) != set(UNNORM_KEYS): + raise StarVLAError( + f"unexpected Qwen2.5 PI normalization profiles: {sorted(norm_stats)}" + ) + for key in UNNORM_KEYS: + normalization_contract(norm_stats, key) + return { + "catalog": catalog, + "catalog_path": catalog_path.resolve(), + "variant": variant, + "qwen": qwen, + "policy_dir": policy_dir.resolve(), + "qwen_dir": qwen_dir.resolve(), + "checkpoint": checkpoint_path.resolve(), + "checkpoint_ready": checkpoint_ready, + "config_yaml": config_yaml.resolve(), + "dataset_statistics": dataset_statistics.resolve(), + "norm_stats": norm_stats, + "config": config, + "source_dir": source_dir, + "source_revision": revision, + "source_probe": source_probe, + "tokenizer": tokenizer, + } + + +def _extract_legacy_git_archive(archive: bytes, destination: Path) -> None: + """Extract verified Git files while ignoring repository-local links.""" + + def regular_file_filter( + member: tarfile.TarInfo, target: str + ) -> tarfile.TarInfo | None: + # The historical tree contains dataset links to machine-local absolute + # paths. They are irrelevant to inference and must never be followed or + # materialized while constructing the verified runtime source tree. + if member.issym() or member.islnk(): + return None + return tarfile.data_filter(member, target) + + with tarfile.open(fileobj=io.BytesIO(archive), mode="r:") as stream: + stream.extractall(destination, filter=regular_file_filter) + + +@contextlib.contextmanager +def _legacy_source_checkout(source_dir: Path): + """Extract the exact historical source tree without mutating git worktrees.""" + + archive = _run_git( + source_dir, + "archive", + "--format=tar", + LEGACY_IMPLEMENTATION_REVISION, + binary=True, + ) + assert isinstance(archive, bytes) + with tempfile.TemporaryDirectory(prefix="starvla-qwen25-pi-legacy-") as temporary: + root = Path(temporary) + _extract_legacy_git_archive(archive, root) + for relative, expected in LEGACY_SOURCE_FILES.items(): + path = root / relative + _ensure_regular_file(path, label=f"extracted legacy source {relative}") + if sha256_file(path) != expected: + raise StarVLAError( + f"extracted legacy source changed unexpectedly: {relative}" + ) + yield root + + +@contextlib.contextmanager +def _config_only_qwen25_bootstrap( + torch: Any, transformers: Any, qwen_dir: Path +): + """Build the Qwen topology in BF16 without loading absent base weights.""" + + model_class = transformers.Qwen2_5_VLForConditionalGeneration + had_override = "from_pretrained" in model_class.__dict__ + original_override = model_class.__dict__.get("from_pretrained") + + def from_config_only(model_id: str | os.PathLike[str], **kwargs: Any): + actual = Path(model_id).resolve() + if actual != qwen_dir.resolve(): + raise StarVLAError( + f"legacy PI wrapper requested unexpected Qwen source: {actual}" + ) + if kwargs.get("torch_dtype") not in (None, "auto", torch.bfloat16): + raise StarVLAError( + f"unexpected Qwen bootstrap dtype: {kwargs.get('torch_dtype')!r}" + ) + config = transformers.AutoConfig.from_pretrained( + actual, local_files_only=True, trust_remote_code=False + ) + declared_model_type = getattr(type(config), "model_type", None) + runtime_model_type = getattr(config, "model_type", None) + text_config = getattr(config, "text_config", config) + vision_config = getattr(config, "vision_config", None) + config_contract = { + "declared_model_type": declared_model_type, + "runtime_model_type": runtime_model_type, + "hidden_size": getattr(text_config, "hidden_size", None), + "layer_count": getattr(text_config, "num_hidden_layers", None), + "vocab_size": getattr(text_config, "vocab_size", None), + "vision_hidden_size": getattr(vision_config, "hidden_size", 1280), + "vision_depth": getattr(vision_config, "depth", 32), + "vision_output_size": getattr(vision_config, "out_hidden_size", 2048), + } + expected_contract = { + "declared_model_type": BACKBONE, + "runtime_model_type": runtime_model_type, + "hidden_size": 2048, + "layer_count": 36, + "vocab_size": 153713, + "vision_hidden_size": 1280, + "vision_depth": 32, + "vision_output_size": 2048, + } + if ( + runtime_model_type not in {BACKBONE, "qwen2_5_vl_text"} + or config_contract != expected_contract + ): + raise StarVLAError( + f"unexpected local Qwen config contract: {config_contract}" + ) + config._attn_implementation = "sdpa" + previous_dtype = torch.get_default_dtype() + try: + torch.set_default_dtype(torch.bfloat16) + with transformers.modeling_utils.no_init_weights(): + return model_class(config) + finally: + torch.set_default_dtype(previous_dtype) + + model_class.from_pretrained = staticmethod(from_config_only) + try: + yield + finally: + if had_override: + model_class.from_pretrained = original_override + else: + delattr(model_class, "from_pretrained") + + +def load_official_framework( + paths: Mapping[str, Any], *, device: str +) -> tuple[Any, dict[str, Any], tempfile.TemporaryDirectory[str]]: + """Load the original checkpoint against its exact historical Python code.""" + + import torch + import transformers + + if not paths["checkpoint_ready"]: + raise StarVLAError( + f"official Qwen2.5 PI checkpoint is absent or incomplete: " + f"{paths['checkpoint']}" + ) + if any(name == "starVLA" or name.startswith("starVLA.") for name in sys.modules): + raise StarVLAError("starVLA was imported before legacy-source verification") + + # Keep the extracted tree alive for as long as the framework class exists. + holder: tempfile.TemporaryDirectory[str] = tempfile.TemporaryDirectory( + prefix="starvla-qwen25-pi-runtime-" + ) + runtime_root = Path(holder.name) + archive = _run_git( + Path(paths["source_dir"]), + "archive", + "--format=tar", + LEGACY_IMPLEMENTATION_REVISION, + binary=True, + ) + assert isinstance(archive, bytes) + _extract_legacy_git_archive(archive, runtime_root) + + sys.path.insert(0, str(runtime_root)) + try: + from starVLA.model.framework import base_framework, share_tools + from starVLA.model.framework import QwenPI + + for module in (base_framework, share_tools, QwenPI): + _assert_module_origin(module, runtime_root) + config = json.loads(json.dumps(paths["config"])) + with _official_qwen25_alias(Path(paths["qwen_dir"])) as qwen_alias: + config["framework"]["qwenvl"]["base_vlm"] = str(qwen_alias) + cfg = share_tools.dict_to_namespace(config) + cfg.trainer.pretrained_checkpoint = None + with _config_only_qwen25_bootstrap( + torch, transformers, Path(paths["qwen_dir"]) + ): + framework = QwenPI.Qwen_PI(cfg) + + try: + state = torch.load( + paths["checkpoint"], + map_location="cpu", + mmap=True, + weights_only=True, + ) + except TypeError: + state = torch.load( + paths["checkpoint"], map_location="cpu", weights_only=True + ) + if not isinstance(state, Mapping) or not state: + raise StarVLAError("official Qwen2.5 PI checkpoint has no state_dict") + framework.load_state_dict(state, strict=True) + del state + gc.collect() + framework.norm_stats = paths["norm_stats"] + + action_model = framework.action_model + if type(framework).__name__ != "Qwen_PI": + raise StarVLAError( + f"unexpected legacy framework class: {type(framework).__name__}" + ) + if len(action_model.model.transformer_blocks) != EXPECTED_DIT_BLOCK_COUNT: + raise StarVLAError( + "legacy PI checkpoint did not instantiate exactly 16 DiT blocks" + ) + if int(action_model.action_horizon) != EXPECTED_ACTION_HORIZON: + raise StarVLAError("legacy PI action horizon changed") + qwen_dtypes = { + parameter.dtype + for parameter in framework.qwen_vl_interface.parameters() + } + policy_dtypes = { + parameter.dtype for parameter in action_model.parameters() + } + if qwen_dtypes != {torch.bfloat16} or policy_dtypes != {torch.float32}: + raise StarVLAError( + "legacy PI dtype boundary changed: " + f"qwen={qwen_dtypes}, policy={policy_dtypes}" + ) + framework = framework.to(device=device, dtype=torch.bfloat16).eval() + runtime_dtypes = {parameter.dtype for parameter in framework.parameters()} + if runtime_dtypes != {torch.bfloat16}: + raise StarVLAError( + f"official --use_bf16 deployment cast failed: {runtime_dtypes}" + ) + return framework, config, holder + except Exception: + holder.cleanup() + raise + finally: + if sys.path and sys.path[0] == str(runtime_root): + del sys.path[0] + + +def _load_images( + image_paths: Iterable[Path], +) -> tuple[list[Any], list[dict[str, Any]]]: + from PIL import Image + + images: list[Any] = [] + records: list[dict[str, Any]] = [] + for path in image_paths: + path = path.resolve() + _ensure_regular_file(path, label="Qwen2.5 PI input image") + try: + with Image.open(path) as opened: + opened.load() + image = opened.convert("RGB") + except (OSError, ValueError) as exc: + raise StarVLAError(f"failed to decode input image {path}: {exc}") from exc + if image.size != (224, 224): + raise StarVLAError( + "Qwen2.5 PI parity requires an already-224x224 image so the " + "released deployment pre-resize is unambiguous" + ) + images.append(image) + records.append( + { + "source_path": str(path), + "source_size": path.stat().st_size, + "source_sha256": sha256_file(path), + "decoded_mode": image.mode, + "decoded_size": list(image.size), + "decoded_pixel_sha256": _image_pixel_sha256(image), + } + ) + if len(images) != 1: + raise StarVLAError("official Qwen2.5 PI oracle requires exactly one image") + return images, records + + +def _render_model_prompt(framework: Any, image: Any, task: str) -> str: + instruction = EXPECTED_COT_TEMPLATE.replace("{instruction}", task) + messages = [{ + "role": "user", + "content": [ + {"type": "image", "image": image}, + {"type": "text", "text": instruction}, + ], + }] + rendered = framework.qwen_vl_interface.processor.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True + ) + if not isinstance(rendered, str): + raise StarVLAError("Qwen2.5 processor returned a non-string prompt") + return rendered + + +def run_official_forward( + framework: Any, + *, + images: Sequence[Any], + task: str, + state: np.ndarray, +) -> dict[str, Any]: + """Run Qwen and the exact legacy 16-block action head with explicit noise.""" + + import torch + + if len(images) != 1: + raise StarVLAError("legacy PI forward requires one image") + if state.shape != (1, 1, EXPECTED_STATE_DIM): + raise StarVLAError(f"legacy PI state must have shape [1,1,7], got {state.shape}") + qwen = framework.qwen_vl_interface + action_model = framework.action_model + qwen_inputs = qwen.build_qwenvl_inputs( + images=[list(images)], instructions=[task] + ) + with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16): + outputs = qwen( + **qwen_inputs, + output_attentions=False, + output_hidden_states=True, + return_dict=True, + ) + hidden = outputs.hidden_states + if hidden is None or len(hidden) != EXPECTED_QWEN_LAYER_COUNT + 1: + raise StarVLAError( + "Qwen2.5 output must expose hidden tuple indices 0..36" + ) + selected = list(hidden[-EXPECTED_DIT_BLOCK_COUNT:]) + if len(selected) != EXPECTED_DIT_BLOCK_COUNT: + raise StarVLAError("legacy PI did not select 16 conditioning states") + if selected[-1].dtype != torch.bfloat16: + raise StarVLAError("Qwen2.5 result_norm boundary must be BF16") + + state_tensor = torch.from_numpy(state).to( + device=selected[-1].device, dtype=torch.bfloat16 + ) + noise = explicit_initial_noise() + noise_tensor = torch.from_numpy(noise).to( + device=selected[-1].device, dtype=torch.bfloat16 + ) + original_randn = torch.randn + noise_calls = 0 + + def explicit_randn(*args: Any, **kwargs: Any): + nonlocal noise_calls + requested_size = kwargs.get("size", args[0] if args else None) + if tuple(requested_size) != tuple(noise_tensor.shape): + raise StarVLAError( + f"legacy PI requested unexpected noise shape: {requested_size}" + ) + if kwargs.get("dtype") != torch.bfloat16: + raise StarVLAError( + f"legacy PI requested unexpected noise dtype: {kwargs.get('dtype')}" + ) + noise_calls += 1 + if noise_calls != 1: + raise StarVLAError("legacy PI requested initial noise more than once") + return noise_tensor.clone() + + torch.randn = explicit_randn + try: + with torch.inference_mode(): + normalized_tensor = action_model.predict_action( + selected, state_tensor + ) + finally: + torch.randn = original_randn + if noise_calls != 1: + raise StarVLAError("legacy PI did not consume the explicit initial noise") + normalized = np.ascontiguousarray( + normalized_tensor.detach().cpu().float().numpy(), dtype=np.float32 + ) + if ( + normalized.shape + != (1, EXPECTED_ACTION_HORIZON, EXPECTED_ACTION_DIM) + or not np.isfinite(normalized).all() + ): + raise StarVLAError( + f"legacy PI produced invalid normalized actions: {normalized.shape}" + ) + qwen_arrays = { + key: value.detach().cpu() + for key, value in qwen_inputs.items() + if isinstance(value, torch.Tensor) + } + required = {"input_ids", "attention_mask", "image_grid_thw"} + if not required.issubset(qwen_arrays): + raise StarVLAError( + f"Qwen2.5 processor inputs are missing: {sorted(required - qwen_arrays.keys())}" + ) + return { + "qwen_inputs": qwen_arrays, + "result_norm": selected[-1].detach(), + "selected_hidden_tuple_indices": EXPECTED_HIDDEN_TUPLE_INDICES, + "policy_conditioning_dtype": "bfloat16", + "initial_noise": noise, + "normalized_actions": normalized, + } + + +def _array_sha256(array: np.ndarray) -> str: + contiguous = np.ascontiguousarray(array) + header = _canonical_json( + {"dtype": contiguous.dtype.str, "shape": list(contiguous.shape)} + ) + return _sha256_bytes(header + b"\x00" + contiguous.tobytes(order="C")) + + +def write_golden( + *, + output_dir: Path, + paths: Mapping[str, Any], + framework: Any, + image_path: Path, + image_record: Mapping[str, Any], + image: Any, + task: str, + unnorm_key: str, + state: np.ndarray, + captures: Mapping[str, Any], + unnormalized: np.ndarray, +) -> Path: + import torch + import transformers + + output_dir = output_dir.resolve() + if output_dir.exists(): + raise StarVLAError(f"golden output directory already exists: {output_dir}") + output_dir.parent.mkdir(parents=True, exist_ok=True) + + noise = np.ascontiguousarray(captures["initial_noise"], dtype=" np.ndarray: + try: + values = [float(item.strip()) for item in value.split(",")] + except ValueError as exc: + raise argparse.ArgumentTypeError( + "--state must contain seven comma-separated finite floats" + ) from exc + if len(values) != EXPECTED_STATE_DIM or not all( + math.isfinite(item) for item in values + ): + raise argparse.ArgumentTypeError( + "--state must contain seven comma-separated finite floats" + ) + return np.asarray(values, dtype=np.float32).reshape(1, 1, EXPECTED_STATE_DIM) + + +def _preflight_record( + paths: Mapping[str, Any], processor: Mapping[str, Any] +) -> dict[str, Any]: + noise = explicit_initial_noise() + return { + "schema_version": SCHEMA_VERSION, + "kind": "starvla_qwen25_pi_preflight", + "variant": VARIANT, + "model_type": MODEL_TYPE, + "backbone": BACKBONE, + "checkpoint": str(paths["checkpoint"]), + "checkpoint_ready": paths["checkpoint_ready"], + "expected_checkpoint": { + "bundle_uuid": official_bundle_uuid( + paths["variant"], paths["catalog"] + ), + "repo_id": OFFICIAL_CHECKPOINT_REPO_ID, + "revision": OFFICIAL_CHECKPOINT_REVISION, + "filename": OFFICIAL_CHECKPOINT_FILENAME, + "size": OFFICIAL_CHECKPOINT_SIZE, + "sha256": OFFICIAL_CHECKPOINT_SHA256, + }, + "qwen": {**paths["tokenizer"], "processor": dict(processor)}, + "conditioning": { + "hidden_tuple_indices": EXPECTED_HIDDEN_TUPLE_INDICES, + "terminal_tap": "result_norm", + "hidden_size": EXPECTED_QWEN_HIDDEN_DIM, + "transport": "native_bfloat16", + }, + "policy": { + "block_count": EXPECTED_DIT_BLOCK_COUNT, + "block_mode": "layerwise_cross_attention_every_block", + "released_config_interleave_self_attention": True, + "use_canonical_dit_forward": False, + "reference_execution_precision": + "whole_model_bf16_via_use_bf16", + "state_dim": EXPECTED_STATE_DIM, + "parameter_dtype": "bfloat16", + "runtime_contract_sha256": PI_RUNTIME_CONTRACT_SHA256, + }, + "action": { + "shape": [1, EXPECTED_ACTION_HORIZON, EXPECTED_ACTION_DIM], + "initial_noise_dtype": "bfloat16", + "initial_noise_algorithm": NOISE_ALGORITHM, + "initial_noise_array_sha256": _array_sha256(noise), + "timestep_ids": EXPECTED_TIMESTEP_IDS, + }, + "image": { + "required_parity_fixture_size": [224, 224], + "reason": + "avoid ambiguity in the released deployment pre-resize path", + }, + "action_gate": { + "reference": "local_original_checkpoint_python", + "metric": "full_tensor_global_relative_l2", + "operator": "<=", + "limit": ACTION_RELATIVE_L2_LIMIT, + "required_outputs": ["normalized_actions", "unnormalized_actions"], + }, + "source_probe": paths["source_probe"], + "effective_config_valid": True, + "golden_created": False, + } + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--checkpoint-root", type=Path, default=Path("ckpts/starvla")) + parser.add_argument("--checkpoint", type=Path) + parser.add_argument("--qwen-model", type=Path) + parser.add_argument( + "--starvla-source", + type=Path, + default=Path("ckpts/starvla/source/starvla"), + ) + parser.add_argument("--catalog", type=Path, default=DEFAULT_CATALOG) + parser.add_argument("--image", action="append", default=[], type=Path) + parser.add_argument( + "--task", default="put small spoon from basket to tray" + ) + parser.add_argument("--unnorm-key", choices=UNNORM_KEYS, default="oxe_bridge") + parser.add_argument( + "--state", + type=_parse_state, + default=_parse_state("0,0,0,0,0,0,0"), + help="seven comma-separated model-space state values", + ) + parser.add_argument("--device", default="cuda:0") + parser.add_argument( + "--output-dir", + type=Path, + default=Path( + "goldens/starvla/qwen25_pi/" + "bridge-episode-000000-frame000-put-spoon" + ), + ) + parser.add_argument("--preflight", "--preflight-only", action="store_true") + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + args = build_parser().parse_args(argv) + holder: tempfile.TemporaryDirectory[str] | None = None + try: + _require_isolated_python() + import torch + import transformers + + validate_runtime_versions( + torch_version=torch.__version__, + torchvision_version=_distribution_version("torchvision"), + transformers_version=transformers.__version__, + numpy_version=np.__version__, + ) + if _distribution_version("qwen-vl-utils") != EXPECTED_QWEN_VL_UTILS_VERSION: + raise StarVLAError( + "qwen-vl-utils must be " + f"{EXPECTED_QWEN_VL_UTILS_VERSION} for the official oracle" + ) + _configure_determinism(torch, seed=0, device=args.device) + paths = validate_local_inputs( + checkpoint_root=args.checkpoint_root, + checkpoint=args.checkpoint, + qwen_model=args.qwen_model, + source_dir=args.starvla_source, + catalog_path=args.catalog, + ) + processor = validate_processor_contract(Path(paths["qwen_dir"])) + if args.preflight: + print( + json.dumps( + _preflight_record(paths, processor), + allow_nan=False, + indent=2, + sort_keys=True, + ) + ) + return 0 + if not paths["checkpoint_ready"]: + raise StarVLAError( + f"official Qwen2.5 PI checkpoint is not ready: {paths['checkpoint']}" + ) + if len(args.image) != 1: + raise StarVLAError("exactly one --image is required") + images, image_records = _load_images(args.image) + framework, _config, holder = load_official_framework( + paths, device=args.device + ) + captures = run_official_forward( + framework, + images=images, + task=args.task, + state=args.state, + ) + unnormalized = unnormalize_actions( + captures["normalized_actions"], paths["norm_stats"], args.unnorm_key + ) + manifest = write_golden( + output_dir=args.output_dir, + paths=paths, + framework=framework, + image_path=args.image[0].resolve(), + image_record=image_records[0], + image=images[0], + task=args.task, + unnorm_key=args.unnorm_key, + state=args.state, + captures=captures, + unnormalized=unnormalized, + ) + print(f"Wrote StarVLA Qwen2.5 PI local-Python golden: {manifest}") + return 0 + except (StarVLAError, OSError, RuntimeError, ValueError) as exc: + print(f"error: {exc}", file=sys.stderr) + return 2 + finally: + if holder is not None: + holder.cleanup() + + +if __name__ == "__main__": + raise SystemExit(main()) From fa9b72b2943159b0b42116e79d604f6041988020 Mon Sep 17 00:00:00 2001 From: JJJYmmm <1650675829@qq.com> Date: Mon, 10 Aug 2026 12:42:56 +0800 Subject: [PATCH 06/11] starvla: add Qwen3-VL GR00T policy --- src/models/starvla/groot_policy.cpp | 855 ++++++++++++ src/models/starvla/groot_policy.h | 82 ++ src/models/starvla/groot_prompt.cpp | 75 ++ src/models/starvla/groot_prompt.h | 17 + tests/starvla/groot_prompt_test.cpp | 86 ++ .../starvla/generate_starvla_groot_golden.py | 1169 +++++++++++++++++ .../starvla/groot_golden_constraints.txt | 13 + .../starvla/serve_starvla_groot_reference.py | 535 ++++++++ 8 files changed, 2832 insertions(+) create mode 100644 src/models/starvla/groot_policy.cpp create mode 100644 src/models/starvla/groot_policy.h create mode 100644 src/models/starvla/groot_prompt.cpp create mode 100644 src/models/starvla/groot_prompt.h create mode 100644 tests/starvla/groot_prompt_test.cpp create mode 100644 tools/hf2gguf/starvla/generate_starvla_groot_golden.py create mode 100644 tools/hf2gguf/starvla/groot_golden_constraints.txt create mode 100644 tools/hf2gguf/starvla/serve_starvla_groot_reference.py diff --git a/src/models/starvla/groot_policy.cpp b/src/models/starvla/groot_policy.cpp new file mode 100644 index 0000000..d9c1296 --- /dev/null +++ b/src/models/starvla/groot_policy.cpp @@ -0,0 +1,855 @@ +#include "models/starvla/groot_policy.h" + +#include "ggml-backend.h" +#include "ggml.h" +#include "gguf.h" +#include "models/ggml_backend.h" +#include "models/gguf_loader.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace robotcpp::starvla { + +namespace { + +constexpr size_t kGraphSize = 16384; +constexpr int kKQMaskPad = 32; + +struct GR00TBlockWeights { + ggml_tensor * ada_norm_weight = nullptr; + ggml_tensor * ada_norm_bias = nullptr; + ggml_tensor * query_weight = nullptr; + ggml_tensor * query_bias = nullptr; + ggml_tensor * key_weight = nullptr; + ggml_tensor * key_bias = nullptr; + ggml_tensor * value_weight = nullptr; + ggml_tensor * value_bias = nullptr; + ggml_tensor * attention_output_weight = nullptr; + ggml_tensor * attention_output_bias = nullptr; + ggml_tensor * feed_forward_input_weight = nullptr; + ggml_tensor * feed_forward_input_bias = nullptr; + ggml_tensor * feed_forward_output_weight = nullptr; + ggml_tensor * feed_forward_output_bias = nullptr; +}; + +struct GR00TWeights { + ggml_tensor * timestep_input_weight = nullptr; + ggml_tensor * timestep_input_bias = nullptr; + ggml_tensor * timestep_output_weight = nullptr; + ggml_tensor * timestep_output_bias = nullptr; + std::vector blocks; + ggml_tensor * output_modulation_weight = nullptr; + ggml_tensor * output_modulation_bias = nullptr; + ggml_tensor * output_projection_weight = nullptr; + ggml_tensor * output_projection_bias = nullptr; + ggml_tensor * action_input_weight = nullptr; + ggml_tensor * action_input_bias = nullptr; + ggml_tensor * action_time_mix_weight = nullptr; + ggml_tensor * action_time_mix_bias = nullptr; + ggml_tensor * action_output_weight = nullptr; + ggml_tensor * action_output_bias = nullptr; + ggml_tensor * velocity_input_weight = nullptr; + ggml_tensor * velocity_input_bias = nullptr; + ggml_tensor * velocity_output_weight = nullptr; + ggml_tensor * velocity_output_bias = nullptr; + ggml_tensor * future_tokens = nullptr; + ggml_tensor * action_position = nullptr; +}; + +int require_key(gguf_context * gguf, const char * key, gguf_type type) { + const int index = gguf_find_key(gguf, key); + if (index < 0) { + throw std::runtime_error(std::string("missing required StarVLA GGUF metadata: ") + key); + } + if (gguf_get_kv_type(gguf, index) != type) { + throw std::runtime_error(std::string("invalid StarVLA GGUF metadata type: ") + key); + } + return index; +} + +std::string require_string(gguf_context * gguf, const char * key) { + return gguf_get_val_str(gguf, require_key(gguf, key, GGUF_TYPE_STRING)); +} + +int require_i32(gguf_context * gguf, const char * key) { + return gguf_get_val_i32(gguf, require_key(gguf, key, GGUF_TYPE_INT32)); +} + +float require_f32(gguf_context * gguf, const char * key) { + return gguf_get_val_f32(gguf, require_key(gguf, key, GGUF_TYPE_FLOAT32)); +} + +bool require_bool(gguf_context * gguf, const char * key) { + return gguf_get_val_bool(gguf, require_key(gguf, key, GGUF_TYPE_BOOL)); +} + +int require_array(gguf_context * gguf, const char * key, gguf_type element_type) { + const int index = require_key(gguf, key, GGUF_TYPE_ARRAY); + if (gguf_get_arr_type(gguf, index) != element_type) { + throw std::runtime_error(std::string("invalid StarVLA GGUF array element type: ") + key); + } + return index; +} + +std::vector require_string_array(gguf_context * gguf, const char * key) { + const int index = require_array(gguf, key, GGUF_TYPE_STRING); + const size_t count = gguf_get_arr_n(gguf, index); + std::vector result; + result.reserve(count); + for (size_t i = 0; i < count; ++i) { + result.emplace_back(gguf_get_arr_str(gguf, index, i)); + } + return result; +} + +std::vector require_i32_array(gguf_context * gguf, const char * key) { + const int index = require_array(gguf, key, GGUF_TYPE_INT32); + const size_t count = gguf_get_arr_n(gguf, index); + const auto * data = static_cast(gguf_get_arr_data(gguf, index)); + if (data == nullptr && count != 0) { + throw std::runtime_error(std::string("missing StarVLA GGUF array data: ") + key); + } + return count == 0 ? std::vector() : std::vector(data, data + count); +} + +std::vector require_f32_array(gguf_context * gguf, const char * key) { + const int index = require_array(gguf, key, GGUF_TYPE_FLOAT32); + const size_t count = gguf_get_arr_n(gguf, index); + const auto * data = static_cast(gguf_get_arr_data(gguf, index)); + if (data == nullptr && count != 0) { + throw std::runtime_error(std::string("missing StarVLA GGUF array data: ") + key); + } + return count == 0 ? std::vector() : std::vector(data, data + count); +} + +std::vector require_bool_array(gguf_context * gguf, const char * key) { + const int index = require_array(gguf, key, GGUF_TYPE_BOOL); + const size_t count = gguf_get_arr_n(gguf, index); + const auto * data = static_cast(gguf_get_arr_data(gguf, index)); + if (data == nullptr && count != 0) { + throw std::runtime_error(std::string("missing StarVLA GGUF array data: ") + key); + } + std::vector result(count); + for (size_t i = 0; i < count; ++i) { + result[i] = data[i] != 0 ? 1 : 0; + } + return result; +} + +std::string profile_key(int profile_index, const char * suffix) { + return "starvla.normalization.profile." + std::to_string(profile_index) + "." + suffix; +} + +bool has_shape(const ggml_tensor * tensor, std::initializer_list expected) { + if (tensor == nullptr || static_cast(ggml_n_dims(tensor)) != expected.size()) { + return false; + } + size_t dimension = 0; + for (const int64_t value : expected) { + if (tensor->ne[dimension++] != value) { + return false; + } + } + return true; +} + +const char * mode_name(backend_mode mode) { + switch (mode) { + case backend_mode::cpu: + return "cpu"; + case backend_mode::cuda: + return "cuda"; + case backend_mode::metal: + return "metal"; + } + return "unknown"; +} + +class GR00TGGUFLoader final : public gguf_loader { + public: + GR00TGGUFLoader(GR00TPolicyConfig & config, GR00TWeights & weights) : config_(config), weights_(weights) {} + + protected: + bool parse_metadata(gguf_context * gguf) override { + if (require_string(gguf, "general.architecture") != "starvla-policy") { + throw std::runtime_error("StarVLA GR00T policy has incompatible general.architecture"); + } + if (require_i32(gguf, "starvla.schema_version") != 1 || + require_string(gguf, "starvla.framework") != "groot") { + throw std::runtime_error("StarVLA policy GGUF is not a supported Qwen GR00T schema"); + } + config_.backbone_arch = require_string(gguf, "starvla.backbone.arch"); + if (config_.backbone_arch != "qwen3_vl" && + config_.backbone_arch != "qwen2_5_vl") { + throw std::runtime_error( + "StarVLA GR00T policy has an unsupported Qwen backbone"); + } + + config_.bundle_uuid = require_string(gguf, "starvla.bundle.uuid"); + if (config_.bundle_uuid.empty()) { + throw std::runtime_error("StarVLA GR00T bundle UUID is missing"); + } + config_.text_filename = require_string(gguf, "starvla.component.text.filename"); + config_.mmproj_filename = require_string(gguf, "starvla.component.mmproj.filename"); + if (config_.text_filename.empty() || config_.mmproj_filename.empty()) { + throw std::runtime_error("StarVLA GR00T component filenames must be non-empty"); + } + config_.qwen_hidden_dim = require_i32(gguf, "starvla.qwen.hidden_size"); + config_.qwen_input_embedding_dim = + require_i32(gguf, "starvla.qwen.input_embedding_size"); + config_.qwen_vocab_size = require_i32(gguf, "starvla.qwen.vocab_size"); + config_.cot_template = require_string(gguf, "starvla.prompt.cot_template"); + const bool qwen25 = config_.backbone_arch == "qwen2_5_vl"; + if (config_.cot_template.empty()) { + throw std::runtime_error("StarVLA GR00T prompt template is missing"); + } + + config_.image_count = require_i32(gguf, "starvla.image.count"); + config_.image_names = require_string_array(gguf, "starvla.image.names"); + config_.image_processor_min_pixels = + require_i32(gguf, "starvla.image.processor_min_pixels"); + config_.image_processor_max_pixels = + require_i32(gguf, "starvla.image.processor_max_pixels"); + config_.image_patch_size = require_i32(gguf, "starvla.image.patch_size"); + config_.image_spatial_merge_size = + require_i32(gguf, "starvla.image.spatial_merge_size"); + config_.image_min_token_count = + require_i32(gguf, "starvla.image.min_token_count"); + config_.image_max_token_count = + require_i32(gguf, "starvla.image.max_token_count"); + config_.dit_width = require_i32(gguf, "starvla.groot.dit_width"); + config_.block_count = require_i32(gguf, "starvla.groot.block_count"); + config_.attention_head_count = require_i32(gguf, "starvla.groot.attention_head_count"); + config_.attention_head_dim = require_i32(gguf, "starvla.groot.attention_head_dim"); + config_.cross_attention_dim = require_i32(gguf, "starvla.groot.cross_attention_dim"); + config_.feed_forward_dim = require_i32(gguf, "starvla.groot.feed_forward_dim"); + config_.output_dim = require_i32(gguf, "starvla.groot.output_dimension"); + config_.mlp_hidden_dim = require_i32(gguf, "starvla.groot.mlp_hidden_dimension"); + config_.future_token_count = require_i32(gguf, "starvla.groot.future_token_count"); + config_.action_position_count = require_i32(gguf, "starvla.groot.action_position_count"); + config_.no_state_sequence_length = require_i32(gguf, "starvla.groot.no_state_sequence_length"); + config_.timestep_projection_dim = require_i32(gguf, "starvla.groot.timestep_projection_dim"); + config_.ada_norm_epsilon = require_f32(gguf, "starvla.groot.ada_norm_epsilon"); + config_.output_norm_epsilon = require_f32(gguf, "starvla.groot.output_norm_epsilon"); + config_.euler_dt = require_f32(gguf, "starvla.groot.euler_dt"); + config_.timestep_ids = require_i32_array(gguf, "starvla.groot.timestep_ids"); + config_.action_dim = require_i32(gguf, "starvla.action.dimension"); + config_.horizon = require_i32(gguf, "starvla.action.horizon"); + + const int64_t expected_input_embedding_dim = + qwen25 ? static_cast(config_.qwen_hidden_dim) + : 4LL * config_.qwen_hidden_dim; + const bool dimensions_valid = + config_.qwen_hidden_dim > 0 && config_.qwen_vocab_size > 0 && config_.dit_width > 0 && + config_.qwen_input_embedding_dim == expected_input_embedding_dim && + config_.dit_width % 2 == 0 && + config_.block_count > 0 && config_.block_count % 2 == 0 && config_.attention_head_count > 0 && + config_.attention_head_dim > 0 && + config_.attention_head_count * config_.attention_head_dim == config_.dit_width && + config_.cross_attention_dim == config_.qwen_hidden_dim && config_.feed_forward_dim > 0 && + config_.output_dim > 0 && config_.mlp_hidden_dim > 0 && + config_.action_dim > 0 && config_.horizon > 0 && config_.future_token_count > 0 && + config_.action_position_count >= config_.horizon && + config_.no_state_sequence_length == config_.future_token_count + config_.horizon && + config_.timestep_projection_dim >= 4 && config_.timestep_projection_dim % 2 == 0 && + std::isfinite(config_.ada_norm_epsilon) && config_.ada_norm_epsilon > 0.0f && + std::isfinite(config_.output_norm_epsilon) && config_.output_norm_epsilon > 0.0f && + std::isfinite(config_.euler_dt) && config_.euler_dt > 0.0f && + config_.timestep_ids.size() == 4 && + config_.image_count > 0 && + config_.image_names.size() == static_cast(config_.image_count) && + config_.image_processor_min_pixels > 0 && + config_.image_processor_max_pixels >= config_.image_processor_min_pixels && + config_.image_patch_size > 0 && config_.image_spatial_merge_size > 0 && + config_.image_min_token_count > 0 && + config_.image_max_token_count >= config_.image_min_token_count; + if (!dimensions_valid) { + throw std::runtime_error("StarVLA GR00T policy metadata has incompatible dimensions"); + } + NormalizationConfig & normalization = config_.normalization; + normalization.clip_actions = require_bool(gguf, "starvla.normalization.clip_actions"); + normalization.binary_threshold = require_f32(gguf, "starvla.normalization.binary_threshold"); + normalization.binary_comparison = require_string(gguf, "starvla.normalization.binary_comparison"); + normalization.continuous_dimensions = + require_i32_array(gguf, "starvla.action.continuous_dimensions"); + normalization.binary_dimensions = require_i32_array(gguf, "starvla.action.binary_dimensions"); + const int profile_count = require_i32(gguf, "starvla.normalization.profile_count"); + const std::vector keys = + require_string_array(gguf, "starvla.normalization.profile_keys"); + if (profile_count <= 0 || keys.size() != static_cast(profile_count)) { + throw std::runtime_error("StarVLA GR00T normalization profile count is inconsistent"); + } + normalization.profiles.clear(); + for (int profile_index = 0; profile_index < profile_count; ++profile_index) { + NormalizationProfile profile; + profile.key = require_string(gguf, profile_key(profile_index, "key").c_str()); + profile.action_q01 = + require_f32_array(gguf, profile_key(profile_index, "action_q01").c_str()); + profile.action_q99 = + require_f32_array(gguf, profile_key(profile_index, "action_q99").c_str()); + profile.action_mask = + require_bool_array(gguf, profile_key(profile_index, "action_mask").c_str()); + if (profile.key != keys[static_cast(profile_index)]) { + throw std::runtime_error("StarVLA GR00T normalization profile order is inconsistent"); + } + normalization.profiles.push_back(std::move(profile)); + } + std::string normalization_error; + if (!validate_normalization_config(normalization, config_.action_dim, normalization_error)) { + throw std::runtime_error(normalization_error); + } + return true; + } + + bool bind_tensors(ggml_context * ctx_data) override { + auto bind = [&](ggml_tensor *& destination, const std::string & name) { + destination = require_tensor(ctx_data, name); + }; + bind(weights_.timestep_input_weight, "starvla.policy.groot.timestep.input.weight"); + bind(weights_.timestep_input_bias, "starvla.policy.groot.timestep.input.bias"); + bind(weights_.timestep_output_weight, "starvla.policy.groot.timestep.output.weight"); + bind(weights_.timestep_output_bias, "starvla.policy.groot.timestep.output.bias"); + weights_.blocks.clear(); + weights_.blocks.reserve(static_cast(config_.block_count)); + for (int block = 0; block < config_.block_count; ++block) { + const std::string prefix = "starvla.policy.groot.block." + std::to_string(block) + "."; + GR00TBlockWeights current; + bind(current.ada_norm_weight, prefix + "ada_norm.weight"); + bind(current.ada_norm_bias, prefix + "ada_norm.bias"); + bind(current.query_weight, prefix + "attention.query.weight"); + bind(current.query_bias, prefix + "attention.query.bias"); + bind(current.key_weight, prefix + "attention.key.weight"); + bind(current.key_bias, prefix + "attention.key.bias"); + bind(current.value_weight, prefix + "attention.value.weight"); + bind(current.value_bias, prefix + "attention.value.bias"); + bind(current.attention_output_weight, prefix + "attention.output.weight"); + bind(current.attention_output_bias, prefix + "attention.output.bias"); + bind(current.feed_forward_input_weight, prefix + "feed_forward.input.weight"); + bind(current.feed_forward_input_bias, prefix + "feed_forward.input.bias"); + bind(current.feed_forward_output_weight, prefix + "feed_forward.output.weight"); + bind(current.feed_forward_output_bias, prefix + "feed_forward.output.bias"); + weights_.blocks.push_back(current); + } + bind(weights_.output_modulation_weight, "starvla.policy.groot.output.modulation.weight"); + bind(weights_.output_modulation_bias, "starvla.policy.groot.output.modulation.bias"); + bind(weights_.output_projection_weight, "starvla.policy.groot.output.projection.weight"); + bind(weights_.output_projection_bias, "starvla.policy.groot.output.projection.bias"); + bind(weights_.action_input_weight, "starvla.policy.groot.action.input.weight"); + bind(weights_.action_input_bias, "starvla.policy.groot.action.input.bias"); + bind(weights_.action_time_mix_weight, "starvla.policy.groot.action.time_mix.weight"); + bind(weights_.action_time_mix_bias, "starvla.policy.groot.action.time_mix.bias"); + bind(weights_.action_output_weight, "starvla.policy.groot.action.output.weight"); + bind(weights_.action_output_bias, "starvla.policy.groot.action.output.bias"); + bind(weights_.velocity_input_weight, "starvla.policy.groot.velocity.input.weight"); + bind(weights_.velocity_input_bias, "starvla.policy.groot.velocity.input.bias"); + bind(weights_.velocity_output_weight, "starvla.policy.groot.velocity.output.weight"); + bind(weights_.velocity_output_bias, "starvla.policy.groot.velocity.output.bias"); + bind(weights_.future_tokens, "starvla.policy.groot.future_tokens.weight"); + bind(weights_.action_position, "starvla.policy.groot.action_position.weight"); + + const int width = config_.dit_width; + if (!has_shape(weights_.timestep_input_weight, {config_.timestep_projection_dim, width}) || + !has_shape(weights_.timestep_input_bias, {width}) || + !has_shape(weights_.timestep_output_weight, {width, width}) || + !has_shape(weights_.timestep_output_bias, {width}) || + !has_shape(weights_.output_modulation_weight, {width, 2 * width}) || + !has_shape(weights_.output_modulation_bias, {2 * width}) || + !has_shape(weights_.output_projection_weight, {width, config_.output_dim}) || + !has_shape(weights_.output_projection_bias, {config_.output_dim}) || + !has_shape(weights_.action_input_weight, {config_.action_dim, width}) || + !has_shape(weights_.action_input_bias, {width}) || + !has_shape(weights_.action_time_mix_weight, {2 * width, width}) || + !has_shape(weights_.action_time_mix_bias, {width}) || + !has_shape(weights_.action_output_weight, {width, width}) || + !has_shape(weights_.action_output_bias, {width}) || + !has_shape(weights_.velocity_input_weight, {config_.output_dim, config_.mlp_hidden_dim}) || + !has_shape(weights_.velocity_input_bias, {config_.mlp_hidden_dim}) || + !has_shape(weights_.velocity_output_weight, {config_.mlp_hidden_dim, config_.action_dim}) || + !has_shape(weights_.velocity_output_bias, {config_.action_dim}) || + !has_shape(weights_.future_tokens, {width, config_.future_token_count}) || + !has_shape(weights_.action_position, {width, config_.action_position_count})) { + throw std::runtime_error("StarVLA GR00T non-block tensor has an incompatible ggml shape"); + } + for (int block = 0; block < config_.block_count; ++block) { + const GR00TBlockWeights & current = weights_.blocks[static_cast(block)]; + const int kv_input_dim = block % 2 == 0 ? config_.cross_attention_dim : width; + if (!has_shape(current.ada_norm_weight, {width, 2 * width}) || + !has_shape(current.ada_norm_bias, {2 * width}) || + !has_shape(current.query_weight, {width, width}) || + !has_shape(current.query_bias, {width}) || + !has_shape(current.key_weight, {kv_input_dim, width}) || + !has_shape(current.key_bias, {width}) || + !has_shape(current.value_weight, {kv_input_dim, width}) || + !has_shape(current.value_bias, {width}) || + !has_shape(current.attention_output_weight, {width, width}) || + !has_shape(current.attention_output_bias, {width}) || + !has_shape(current.feed_forward_input_weight, {width, config_.feed_forward_dim}) || + !has_shape(current.feed_forward_input_bias, {config_.feed_forward_dim}) || + !has_shape(current.feed_forward_output_weight, {config_.feed_forward_dim, width}) || + !has_shape(current.feed_forward_output_bias, {width})) { + throw std::runtime_error("StarVLA GR00T transformer block tensor has an incompatible ggml shape"); + } + } + return true; + } + + private: + GR00TPolicyConfig & config_; + GR00TWeights & weights_; +}; + +std::vector timestep_projection_table(const GR00TPolicyConfig & config) { + const int dim = config.timestep_projection_dim; + const int half = dim / 2; + const float denominator = static_cast(half - 1); + std::vector table(static_cast(dim) * 4, 0.0f); + for (int step = 0; step < 4; ++step) { + const float timestep = static_cast(config.timestep_ids[static_cast(step)]); + float * row = table.data() + static_cast(step) * dim; + for (int index = 0; index < half; ++index) { + const float frequency = std::exp(-std::log(10000.0f) * static_cast(index) / denominator); + const float angle = timestep * frequency; + row[index] = std::cos(angle); + row[index + half] = std::sin(angle); + } + } + return table; +} + +std::vector action_time_table(const GR00TPolicyConfig & config) { + const int dim = config.dit_width; + const int half = dim / 2; + const float denominator = static_cast(half); + std::vector table(static_cast(dim) * 4, 0.0f); + for (int step = 0; step < 4; ++step) { + const float timestep = static_cast(config.timestep_ids[static_cast(step)]); + float * row = table.data() + static_cast(step) * dim; + for (int index = 0; index < half; ++index) { + const float frequency = std::exp(-std::log(10000.0f) * static_cast(index) / denominator); + const float angle = timestep * frequency; + row[index] = std::sin(angle); + row[index + half] = std::cos(angle); + } + } + return table; +} + +} // namespace + +struct GR00TPolicy::Impl { + GR00TPolicyConfig config; + GR00TWeights weights; + gguf_load_result loaded; + ggml_backend_t backend_cpu = nullptr; + std::vector backends; + ggml_backend_sched_t scheduler = nullptr; + backend_buft_policy buft_policy; + backend_mode mode = backend_mode::cpu; + int n_threads = 0; + int verbosity = 0; + + ggml_context * graph_context = nullptr; + ggml_cgraph * graph = nullptr; + ggml_tensor * hidden_input = nullptr; + ggml_tensor * cross_mask_input = nullptr; + ggml_tensor * noise_input = nullptr; + ggml_tensor * timestep_projection_input = nullptr; + ggml_tensor * action_time_input = nullptr; + ggml_tensor * scalar_one_input = nullptr; + ggml_tensor * output = nullptr; + size_t conditioning_token_count = 0; + std::vector timestep_table; + std::vector action_table; + + ~Impl() { + clear_graph(); + if (scheduler != nullptr) { + ggml_backend_sched_synchronize(scheduler); + ggml_backend_sched_free(scheduler); + scheduler = nullptr; + } + if (loaded.model_buffer != nullptr) { + ggml_backend_buffer_free(loaded.model_buffer); + loaded.model_buffer = nullptr; + } + if (loaded.ctx_data != nullptr) { + ggml_free(loaded.ctx_data); + loaded.ctx_data = nullptr; + } + if (loaded.gguf != nullptr) { + gguf_free(loaded.gguf); + loaded.gguf = nullptr; + } + for (ggml_backend_t backend : backends) { + if (backend != nullptr) { + ggml_backend_free(backend); + } + } + backends.clear(); + backend_cpu = nullptr; + } + + void clear_graph() { + if (scheduler != nullptr) { + ggml_backend_sched_synchronize(scheduler); + ggml_backend_sched_reset(scheduler); + } + if (graph_context != nullptr) { + ggml_free(graph_context); + graph_context = nullptr; + } + graph = nullptr; + hidden_input = nullptr; + cross_mask_input = nullptr; + noise_input = nullptr; + timestep_projection_input = nullptr; + action_time_input = nullptr; + scalar_one_input = nullptr; + output = nullptr; + conditioning_token_count = 0; + } + + void build_graph(size_t token_count) { + clear_graph(); + if (token_count == 0 || token_count > static_cast(std::numeric_limits::max())) { + throw std::runtime_error("invalid StarVLA GR00T conditioning token count"); + } + + ggml_init_params params{}; + params.mem_size = kGraphSize * ggml_tensor_overhead() + ggml_graph_overhead_custom(kGraphSize, false); + params.mem_buffer = nullptr; + params.no_alloc = true; + graph_context = ggml_init(params); + if (graph_context == nullptr) { + throw std::runtime_error("failed to initialize StarVLA GR00T graph context"); + } + + const int width = config.dit_width; + const int heads = config.attention_head_count; + const int head_dim = config.attention_head_dim; + const int sequence_length = config.no_state_sequence_length; + const int mask_queries = GGML_PAD(sequence_length, kKQMaskPad); + + hidden_input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, config.qwen_hidden_dim, + static_cast(token_count)); + cross_mask_input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, + static_cast(token_count), mask_queries); + noise_input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, config.action_dim, config.horizon); + timestep_projection_input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, + config.timestep_projection_dim, 4); + action_time_input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, width, 4); + scalar_one_input = ggml_new_tensor_1d(graph_context, GGML_TYPE_F32, 1); + if (hidden_input == nullptr || cross_mask_input == nullptr || noise_input == nullptr || + timestep_projection_input == nullptr || action_time_input == nullptr || scalar_one_input == nullptr) { + throw std::runtime_error("failed to create StarVLA GR00T graph inputs"); + } + ggml_set_name(hidden_input, "starvla_groot_qwen_hidden_states"); + ggml_set_name(cross_mask_input, "starvla_groot_qwen_attention_mask"); + ggml_set_name(noise_input, "starvla_groot_initial_noise"); + ggml_set_name(timestep_projection_input, "starvla_groot_timestep_projection_table"); + ggml_set_name(action_time_input, "starvla_groot_action_time_table"); + ggml_set_name(scalar_one_input, "starvla_groot_scalar_one"); + ggml_set_input(hidden_input); + ggml_set_input(cross_mask_input); + ggml_set_input(noise_input); + ggml_set_input(timestep_projection_input); + ggml_set_input(action_time_input); + ggml_set_input(scalar_one_input); + + auto f32 = [&](ggml_tensor * tensor) { + return tensor->type == GGML_TYPE_F32 ? tensor : ggml_cast(graph_context, tensor, GGML_TYPE_F32); + }; + auto linear = [&](ggml_tensor * value, ggml_tensor * weight, ggml_tensor * bias) { + ggml_tensor * projected = ggml_mul_mat(graph_context, weight, value); + ggml_mul_mat_set_prec(projected, GGML_PREC_F32); + return ggml_add(graph_context, projected, f32(bias)); + }; + auto ada_norm = [&](ggml_tensor * value, ggml_tensor * temb, const GR00TBlockWeights & block) { + ggml_tensor * modulation = linear(ggml_silu(graph_context, temb), block.ada_norm_weight, + block.ada_norm_bias); + ggml_tensor * scale = ggml_view_1d(graph_context, modulation, width, 0); + ggml_tensor * shift = ggml_view_1d(graph_context, modulation, width, + static_cast(width) * sizeof(float)); + ggml_tensor * normalized = ggml_norm(graph_context, value, config.ada_norm_epsilon); + ggml_tensor * one_plus_scale = ggml_add(graph_context, scale, scalar_one_input); + return ggml_add(graph_context, ggml_mul(graph_context, normalized, one_plus_scale), shift); + }; + auto attention = [&](ggml_tensor * query_source, ggml_tensor * key_value_source, + ggml_tensor * mask, const GR00TBlockWeights & block) { + const int64_t query_count = query_source->ne[1]; + const int64_t key_value_count = key_value_source->ne[1]; + ggml_tensor * query = linear(query_source, block.query_weight, block.query_bias); + ggml_tensor * key = linear(key_value_source, block.key_weight, block.key_bias); + ggml_tensor * value = linear(key_value_source, block.value_weight, block.value_bias); + query = ggml_reshape_3d(graph_context, query, head_dim, heads, query_count); + key = ggml_reshape_3d(graph_context, key, head_dim, heads, key_value_count); + value = ggml_reshape_3d(graph_context, value, head_dim, heads, key_value_count); + query = ggml_permute(graph_context, query, 0, 2, 1, 3); + key = ggml_permute(graph_context, key, 0, 2, 1, 3); + value = ggml_cont(graph_context, ggml_permute(graph_context, value, 1, 2, 0, 3)); + ggml_tensor * scores = ggml_mul_mat(graph_context, key, query); + ggml_mul_mat_set_prec(scores, GGML_PREC_F32); + scores = ggml_soft_max_ext(graph_context, scores, mask, + 1.0f / std::sqrt(static_cast(head_dim)), 0.0f); + ggml_tensor * attended = ggml_mul_mat(graph_context, value, scores); + ggml_mul_mat_set_prec(attended, GGML_PREC_F32); + attended = ggml_permute(graph_context, attended, 0, 2, 1, 3); + attended = ggml_cont_2d(graph_context, attended, width, query_count); + return linear(attended, block.attention_output_weight, block.attention_output_bias); + }; + ggml_tensor * future = f32(weights.future_tokens); + ggml_tensor * position_view = ggml_view_2d(graph_context, weights.action_position, width, config.horizon, + weights.action_position->nb[1], 0); + ggml_tensor * position = f32(position_view); + ggml_tensor * actions = noise_input; + + for (int step = 0; step < 4; ++step) { + ggml_tensor * timestep_projection = ggml_view_1d( + graph_context, timestep_projection_input, config.timestep_projection_dim, + static_cast(step) * config.timestep_projection_dim * sizeof(float)); + ggml_tensor * temb = linear(timestep_projection, weights.timestep_input_weight, + weights.timestep_input_bias); + temb = ggml_silu(graph_context, temb); + temb = linear(temb, weights.timestep_output_weight, weights.timestep_output_bias); + + ggml_tensor * action_features = linear(actions, weights.action_input_weight, weights.action_input_bias); + ggml_tensor * action_time = ggml_view_1d( + graph_context, action_time_input, width, + static_cast(step) * width * sizeof(float)); + action_time = ggml_repeat(graph_context, action_time, action_features); + action_features = ggml_concat(graph_context, action_features, action_time, 0); + action_features = linear(action_features, weights.action_time_mix_weight, + weights.action_time_mix_bias); + action_features = ggml_silu(graph_context, action_features); + action_features = linear(action_features, weights.action_output_weight, weights.action_output_bias); + action_features = ggml_add(graph_context, action_features, position); + ggml_tensor * hidden = ggml_concat(graph_context, future, action_features, 1); + for (int block_index = 0; block_index < config.block_count; ++block_index) { + const GR00TBlockWeights & block = weights.blocks[static_cast(block_index)]; + ggml_tensor * normalized = ada_norm(hidden, temb, block); + ggml_tensor * attended = block_index % 2 == 0 + ? attention(normalized, hidden_input, cross_mask_input, block) + : attention(normalized, normalized, nullptr, block); + hidden = ggml_add(graph_context, hidden, attended); + ggml_tensor * ff = ggml_norm(graph_context, hidden, config.ada_norm_epsilon); + ff = linear(ff, block.feed_forward_input_weight, block.feed_forward_input_bias); + ff = ggml_gelu(graph_context, ff); + ff = linear(ff, block.feed_forward_output_weight, block.feed_forward_output_bias); + hidden = ggml_add(graph_context, hidden, ff); + } + + ggml_tensor * output_modulation = linear(ggml_silu(graph_context, temb), + weights.output_modulation_weight, + weights.output_modulation_bias); + // DiT output uses shift then scale, unlike AdaLayerNorm's scale then shift. + ggml_tensor * shift = ggml_view_1d(graph_context, output_modulation, width, 0); + ggml_tensor * scale = ggml_view_1d(graph_context, output_modulation, width, + static_cast(width) * sizeof(float)); + hidden = ggml_norm(graph_context, hidden, config.output_norm_epsilon); + hidden = ggml_mul(graph_context, hidden, ggml_add(graph_context, scale, scalar_one_input)); + hidden = ggml_add(graph_context, hidden, shift); + hidden = linear(hidden, weights.output_projection_weight, weights.output_projection_bias); + hidden = ggml_relu(graph_context, + linear(hidden, weights.velocity_input_weight, weights.velocity_input_bias)); + hidden = linear(hidden, weights.velocity_output_weight, weights.velocity_output_bias); + ggml_tensor * velocity = ggml_view_2d( + graph_context, hidden, config.action_dim, config.horizon, hidden->nb[1], + static_cast(config.future_token_count) * hidden->nb[1]); + actions = ggml_add(graph_context, actions, ggml_scale(graph_context, velocity, config.euler_dt)); + } + + output = actions; + ggml_set_name(output, "starvla_groot_normalized_actions"); + ggml_set_output(output); + graph = ggml_new_graph_custom(graph_context, kGraphSize, false); + if (graph == nullptr) { + throw std::runtime_error("failed to create StarVLA GR00T graph"); + } + ggml_build_forward_expand(graph, output); + ggml_backend_sched_reset(scheduler); + if (!ggml_backend_sched_alloc_graph(scheduler, graph)) { + throw std::runtime_error("failed to allocate StarVLA GR00T graph"); + } + + conditioning_token_count = token_count; + } +}; + +GR00TPolicy::GR00TPolicy(std::unique_ptr impl) : impl_(std::move(impl)) {} + +GR00TPolicy::~GR00TPolicy() = default; + +std::unique_ptr GR00TPolicy::load(const std::string & path, int n_threads, int verbosity, + std::string & error) { + error.clear(); + if (path.empty()) { + error = "StarVLA GR00T policy path is required"; + return nullptr; + } + + std::unique_ptr impl(new Impl()); + impl->n_threads = n_threads; + impl->verbosity = verbosity; + try { + backend_scheduler_config scheduler_config; + scheduler_config.max_nodes = static_cast(kGraphSize); + scheduler_config.parallel = false; + scheduler_config.op_offload = true; + backend_loader backend; + if (!backend.load(impl->backend_cpu, impl->backends, impl->scheduler, impl->buft_policy, true, + scheduler_config, verbosity)) { + error = "failed to initialize StarVLA GR00T backend: " + backend.error(); + return nullptr; + } + impl->mode = backend.mode(); + + GR00TGGUFLoader loader(impl->config, impl->weights); + if (!loader.load(path.c_str(), impl->buft_policy.model_buft, impl->loaded, verbosity)) { + error = loader.error(); + return nullptr; + } + if (impl->loaded.ctx_data == nullptr || impl->loaded.model_buffer == nullptr) { + error = "StarVLA GR00T policy GGUF has no tensors"; + return nullptr; + } + ggml_backend_buffer_set_usage(impl->loaded.model_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + impl->timestep_table = timestep_projection_table(impl->config); + impl->action_table = action_time_table(impl->config); + if (verbosity >= 1) { + std::fprintf(stderr, + "%s: backend=%s qwen=%d width=%d blocks=%d horizon=%d action_dim=%d profiles=%zu\n", + __func__, mode_name(impl->mode), impl->config.qwen_hidden_dim, impl->config.dit_width, + impl->config.block_count, impl->config.horizon, impl->config.action_dim, + impl->config.normalization.profiles.size()); + } + } catch (const std::exception & exception) { + error = exception.what(); + return nullptr; + } + return std::unique_ptr(new GR00TPolicy(std::move(impl))); +} + +bool GR00TPolicy::evaluate(const float * qwen_hidden_states, size_t hidden_element_count, + const uint8_t * qwen_attention_mask, size_t mask_element_count, + const float * initial_noise, size_t noise_element_count, + std::vector & normalized_actions, std::string & error) { + normalized_actions.clear(); + error.clear(); + if (impl_ == nullptr || impl_->scheduler == nullptr) { + error = "StarVLA GR00T policy is not initialized"; + return false; + } + if (qwen_hidden_states == nullptr || qwen_attention_mask == nullptr || initial_noise == nullptr || + mask_element_count == 0 || mask_element_count > static_cast(std::numeric_limits::max()) || + mask_element_count > std::numeric_limits::max() / + static_cast(impl_->config.qwen_hidden_dim) || + hidden_element_count != mask_element_count * static_cast(impl_->config.qwen_hidden_dim)) { + error = "StarVLA GR00T Qwen conditioning tensor or attention mask has an incompatible shape"; + return false; + } + const size_t expected_noise = static_cast(impl_->config.horizon) * impl_->config.action_dim; + if (noise_element_count != expected_noise) { + error = "StarVLA GR00T initial-noise tensor has an incompatible shape"; + return false; + } + if (std::any_of(qwen_hidden_states, qwen_hidden_states + hidden_element_count, + [](float value) { return !std::isfinite(value); }) || + std::any_of(initial_noise, initial_noise + noise_element_count, + [](float value) { return !std::isfinite(value); })) { + error = "StarVLA GR00T conditioning and initial noise must be finite"; + return false; + } + bool has_valid_token = false; + for (size_t token = 0; token < mask_element_count; ++token) { + if (qwen_attention_mask[token] > 1) { + error = "StarVLA GR00T attention mask values must be zero or one"; + return false; + } + has_valid_token = has_valid_token || qwen_attention_mask[token] != 0; + } + if (!has_valid_token) { + error = "StarVLA GR00T attention mask must contain at least one valid token"; + return false; + } + + try { + if (impl_->graph == nullptr || impl_->conditioning_token_count != mask_element_count) { + impl_->build_graph(mask_element_count); + } + } catch (const std::exception & exception) { + error = exception.what(); + return false; + } + + const int query_count = impl_->config.no_state_sequence_length; + const int padded_queries = GGML_PAD(query_count, kKQMaskPad); + std::vector additive_mask(mask_element_count * static_cast(padded_queries), + -std::numeric_limits::infinity()); + for (int query = 0; query < query_count; ++query) { + float * row = additive_mask.data() + static_cast(query) * mask_element_count; + for (size_t token = 0; token < mask_element_count; ++token) { + row[token] = qwen_attention_mask[token] != 0 ? 0.0f : + -std::numeric_limits::infinity(); + } + } + + ggml_backend_tensor_set(impl_->hidden_input, qwen_hidden_states, 0, + hidden_element_count * sizeof(float)); + ggml_backend_tensor_set(impl_->cross_mask_input, additive_mask.data(), 0, + additive_mask.size() * sizeof(float)); + ggml_backend_tensor_set(impl_->noise_input, initial_noise, 0, noise_element_count * sizeof(float)); + ggml_backend_tensor_set(impl_->timestep_projection_input, impl_->timestep_table.data(), 0, + impl_->timestep_table.size() * sizeof(float)); + ggml_backend_tensor_set(impl_->action_time_input, impl_->action_table.data(), 0, + impl_->action_table.size() * sizeof(float)); + const float one = 1.0f; + ggml_backend_tensor_set(impl_->scalar_one_input, &one, 0, sizeof(one)); + set_backend_threads(impl_->backends, impl_->n_threads); + if (ggml_backend_sched_graph_compute(impl_->scheduler, impl_->graph) != GGML_STATUS_SUCCESS) { + error = "StarVLA GR00T graph compute failed"; + return false; + } + + normalized_actions.resize(expected_noise); + ggml_backend_tensor_get(impl_->output, normalized_actions.data(), 0, expected_noise * sizeof(float)); + if (std::any_of(normalized_actions.begin(), normalized_actions.end(), + [](float value) { return !std::isfinite(value); })) { + normalized_actions.clear(); + error = "StarVLA GR00T graph produced non-finite actions"; + return false; + } + return true; +} + +bool GR00TPolicy::unnormalize(const std::vector & normalized_actions, + const std::string & profile_key_value, std::vector & actions, + std::string & error) const { + if (impl_ == nullptr) { + actions.clear(); + error = "StarVLA GR00T policy is not initialized"; + return false; + } + return denormalize_actions(impl_->config.normalization, profile_key_value, normalized_actions, + impl_->config.horizon, impl_->config.action_dim, actions, error); +} + +const GR00TPolicyConfig & GR00TPolicy::config() const { + if (impl_ == nullptr) { + throw std::runtime_error("StarVLA GR00T policy is not initialized"); + } + return impl_->config; +} + +const char * GR00TPolicy::backend_name() const { + return impl_ != nullptr ? mode_name(impl_->mode) : "unknown"; +} + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/groot_policy.h b/src/models/starvla/groot_policy.h new file mode 100644 index 0000000..c505035 --- /dev/null +++ b/src/models/starvla/groot_policy.h @@ -0,0 +1,82 @@ +#pragma once + +#include "models/starvla/normalization.h" + +#include +#include +#include +#include +#include + +namespace robotcpp::starvla { + +struct GR00TPolicyConfig { + std::string backbone_arch; + std::string bundle_uuid; + std::string text_filename; + std::string mmproj_filename; + int qwen_hidden_dim = 0; + int qwen_input_embedding_dim = 0; + int qwen_vocab_size = 0; + std::string cot_template; + int image_count = 0; + std::vector image_names; + int image_processor_min_pixels = 0; + int image_processor_max_pixels = 0; + int image_patch_size = 0; + int image_spatial_merge_size = 0; + int image_min_token_count = 0; + int image_max_token_count = 0; + int dit_width = 0; + int block_count = 0; + int attention_head_count = 0; + int attention_head_dim = 0; + int cross_attention_dim = 0; + int feed_forward_dim = 0; + int output_dim = 0; + int mlp_hidden_dim = 0; + int action_dim = 0; + int horizon = 0; + int future_token_count = 0; + int action_position_count = 0; + int no_state_sequence_length = 0; + int timestep_projection_dim = 0; + float ada_norm_epsilon = 0.0f; + float output_norm_epsilon = 0.0f; + float euler_dt = 0.0f; + std::vector timestep_ids; + NormalizationConfig normalization; +}; + +class GR00TPolicy { + public: + ~GR00TPolicy(); + + GR00TPolicy(const GR00TPolicy &) = delete; + GR00TPolicy & operator=(const GR00TPolicy &) = delete; + + static std::unique_ptr load(const std::string & path, int n_threads, int verbosity, + std::string & error); + + // qwen_hidden_states is token-major [token_count, qwen_hidden_dim]. The mask + // follows torch SDPA semantics: non-zero entries participate in attention. + // initial_noise is token-major [horizon, action_dim]. + bool evaluate(const float * qwen_hidden_states, size_t hidden_element_count, + const uint8_t * qwen_attention_mask, size_t mask_element_count, + const float * initial_noise, size_t noise_element_count, + std::vector & normalized_actions, std::string & error); + bool unnormalize(const std::vector & normalized_actions, const std::string & profile_key, + std::vector & actions, std::string & error) const; + + const GR00TPolicyConfig & config() const; + const char * backend_name() const; + + private: + struct Impl; + + explicit GR00TPolicy(std::unique_ptr impl); + + std::unique_ptr impl_; +}; + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/groot_prompt.cpp b/src/models/starvla/groot_prompt.cpp new file mode 100644 index 0000000..a1bbe78 --- /dev/null +++ b/src/models/starvla/groot_prompt.cpp @@ -0,0 +1,75 @@ +#include "models/starvla/groot_prompt.h" + +#include + +namespace robotcpp::starvla { +namespace { + +constexpr const char * kInstructionPlaceholder = "{instruction}"; +constexpr const char * kMtmdMediaMarker = "<__media__>"; + +bool contains_nul(const std::string & value) { + return value.find('\0') != std::string::npos; +} + +void replace_all(std::string & value, const std::string & needle, + const std::string & replacement) { + size_t offset = 0; + while ((offset = value.find(needle, offset)) != std::string::npos) { + value.replace(offset, needle.size(), replacement); + offset += replacement.size(); + } +} + +bool build_instruction(const char * framework, const std::string & cot_template, + const std::string & task, std::string & instruction, + std::string & error) { + instruction.clear(); + error.clear(); + + if (task.empty()) { + error = std::string("StarVLA ") + framework + " task must not be empty"; + return false; + } + if (contains_nul(task) || contains_nul(cot_template)) { + error = std::string("StarVLA ") + framework + + " prompt contains an embedded NUL byte"; + return false; + } + if (task.find(kMtmdMediaMarker) != std::string::npos || + cot_template.find(kMtmdMediaMarker) != std::string::npos) { + error = std::string("StarVLA ") + framework + + " prompt contains the reserved mtmd media marker"; + return false; + } + if (cot_template.find(kInstructionPlaceholder) == std::string::npos) { + error = std::string("StarVLA ") + framework + + " CoT template is missing {instruction}"; + return false; + } + + std::string wrapped = cot_template; + replace_all(wrapped, kInstructionPlaceholder, task); + instruction = std::move(wrapped); + return true; +} + +} // namespace + +bool build_groot_instruction(const std::string & cot_template, const std::string & task, + std::string & instruction, std::string & error) { + return build_instruction("GR00T", cot_template, task, instruction, error); +} + +bool build_pi_v3_instruction(const std::string & cot_template, const std::string & task, + std::string & instruction, std::string & error) { + return build_instruction("PI_v3", cot_template, task, instruction, error); +} + +bool build_fast_instruction(const std::string & cot_template, + const std::string & task, + std::string & instruction, std::string & error) { + return build_instruction("FAST", cot_template, task, instruction, error); +} + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/groot_prompt.h b/src/models/starvla/groot_prompt.h new file mode 100644 index 0000000..27e2950 --- /dev/null +++ b/src/models/starvla/groot_prompt.h @@ -0,0 +1,17 @@ +#pragma once + +#include + +namespace robotcpp::starvla { + +bool build_groot_instruction(const std::string & cot_template, const std::string & task, + std::string & instruction, std::string & error); + +bool build_pi_v3_instruction(const std::string & cot_template, const std::string & task, + std::string & instruction, std::string & error); + +bool build_fast_instruction(const std::string & cot_template, + const std::string & task, + std::string & instruction, std::string & error); + +} // namespace robotcpp::starvla diff --git a/tests/starvla/groot_prompt_test.cpp b/tests/starvla/groot_prompt_test.cpp new file mode 100644 index 0000000..fe06d0d --- /dev/null +++ b/tests/starvla/groot_prompt_test.cpp @@ -0,0 +1,86 @@ +#include "models/starvla/groot_prompt.h" + +#include +#include +#include + +namespace { + +void require(bool condition, const char * message) { + if (!condition) { + std::cerr << "FAIL: " << message << '\n'; + std::exit(1); + } +} + +void require_rejected(const std::string & cot_template, const std::string & task, + const char * message) { + std::string instruction = "stale instruction"; + std::string error = "stale error"; + require(!robotcpp::starvla::build_groot_instruction(cot_template, task, instruction, error), + message); + require(instruction.empty(), "rejected prompts must clear their instruction output"); + require(!error.empty(), "rejected prompts must explain the contract violation"); +} + +} // namespace + +int main() { + using robotcpp::starvla::build_groot_instruction; + using robotcpp::starvla::build_pi_v3_instruction; + + const std::string official_template = + "Your task is {instruction}. To identify the key objects for your task. " + "Locate their bounding boxes in [x1,y1,x2,y2] format."; + + std::string instruction; + std::string error = "stale error"; + require(build_groot_instruction(official_template, "grab the block", instruction, error), + "official GR00T prompt must build"); + require(instruction == + "Your task is grab the block. To identify the key objects for your task. " + "Locate their bounding boxes in [x1,y1,x2,y2] format.", + "official GR00T prompt must preserve the checkpoint template exactly"); + require(error.empty(), "successful prompt construction must clear stale errors"); + + require(build_pi_v3_instruction(official_template, "grab the block", instruction, error), + "official PI_v3 prompt must build"); + require(instruction == + "Your task is grab the block. To identify the key objects for your task. " + "Locate their bounding boxes in [x1,y1,x2,y2] format.", + "PI_v3 must use the same pinned CoT replacement contract"); + require(!build_pi_v3_instruction(official_template, "", instruction, error), + "empty PI_v3 tasks must be rejected"); + require(error.find("PI_v3") != std::string::npos, + "PI_v3 prompt errors must identify the active framework"); + + require(build_groot_instruction(official_template, "grab the block.", instruction, error), + "punctuated tasks must build"); + require(instruction == + "Your task is grab the block.. To identify the key objects for your task. " + "Locate their bounding boxes in [x1,y1,x2,y2] format.", + "task punctuation must not be normalized"); + + require(build_groot_instruction("First {instruction}; then {instruction}.", "pick", instruction, + error), + "templates with repeated placeholders must build"); + require(instruction == "First pick; then pick.", + "every instruction placeholder must be replaced"); + + require_rejected(official_template, "", "empty tasks must be rejected"); + require_rejected(official_template, std::string("grab\0now", 8), + "embedded NUL bytes in tasks must be rejected"); + require_rejected(std::string("Use {instruction}\0now", 21), "grab", + "embedded NUL bytes in templates must be rejected"); + require_rejected(official_template, "grab <__media__> now", + "reserved mtmd media markers in tasks must be rejected"); + require_rejected("<__media__>{instruction}", "grab", + "reserved mtmd media markers in templates must be rejected"); + require_rejected("Your task is ready.", "grab", + "templates without an instruction placeholder must be rejected"); + require_rejected("Your task is { instruction }.", "grab", + "lookalike placeholders must not satisfy the template contract"); + + std::cout << "starvla GR00T prompt tests passed\n"; + return 0; +} diff --git a/tools/hf2gguf/starvla/generate_starvla_groot_golden.py b/tools/hf2gguf/starvla/generate_starvla_groot_golden.py new file mode 100644 index 0000000..d90ede7 --- /dev/null +++ b/tools/hf2gguf/starvla/generate_starvla_groot_golden.py @@ -0,0 +1,1169 @@ +#!/usr/bin/env python3 +"""Generate an independent, fixed-noise oracle for official Qwen-GR00T.""" + +from __future__ import annotations + +import argparse +import datetime as dt +import gc +import hashlib +import importlib.metadata +import json +import os +import platform +import random +import shutil +import sys +import tempfile +from pathlib import Path +from typing import Any, Iterable, Mapping, Sequence + +import numpy as np + + +TOOLS_DIR = Path(__file__).resolve().parent +if str(TOOLS_DIR) not in sys.path: + sys.path.insert(0, str(TOOLS_DIR)) + +from generate_starvla_pi_v3_golden import ( # noqa: E402 + EXPECTED_ACCELERATE_VERSION, + EXPECTED_DIFFUSERS_VERSION, + EXPECTED_NUMPY_VERSION, + EXPECTED_OMEGACONF_VERSION, + EXPECTED_PILLOW_VERSION, + EXPECTED_QWEN2VL_IMAGE_PROCESSING_FAST_SHA256, + EXPECTED_QWEN2VL_IMAGE_PROCESSING_SHA256, + EXPECTED_QWEN3VL_MODELING_SHA256, + EXPECTED_QWEN3VL_PROCESSING_SHA256, + EXPECTED_SAFETENSORS_VERSION, + EXPECTED_TOKENIZERS_VERSION, + EXPECTED_TORCHVISION_VERSION, + EXPECTED_TORCH_VERSION, + EXPECTED_TRANSFORMERS_GENERIC_SHA256, + EXPECTED_TRANSFORMERS_VERSION, + OFFICIAL_ENVIRONMENT_FREEZE, + _array_record, + _array_sha256, + _assert_module_origin, + _canonical_json, + _config_only_qwen_bootstrap, + _configure_determinism, + _official_qwen_model_alias, + _sha256_bytes, + validate_runtime_versions, + verify_pinned_source_checkout, + verify_transformers_qwen3vl_recorder_semantics, +) +from starvla_checkpoint import ( # noqa: E402 + DEFAULT_CATALOG, + StarVLAError, + get_variant, + load_catalog, + official_bundle_uuid, + resolve_effective_config, + sha256_file, + verify_catalog_files, + verify_checkpoint_file, +) + + +GOLDEN_SCHEMA_VERSION = 1 +GOLDEN_KIND = "starvla_groot_official_python_oracle" +RUNNER_CONTRACT_KIND = "starvla_groot_runner_contract" +SUPPORTED_VARIANT = "groot" +SEED = 0 +EXPECTED_ACTION_HORIZON = 16 +EXPECTED_ACTION_DIM = 7 +EXPECTED_QWEN_HIDDEN_DIM = 2560 +EXPECTED_DIT_WIDTH = 768 +EXPECTED_DIT_OUTPUT_DIM = 1024 +EXPECTED_DIT_BLOCK_COUNT = 16 +EXPECTED_FUTURE_TOKEN_COUNT = 32 +EXPECTED_SEQUENCE_LENGTH = 48 +EXPECTED_TIMESTEP_IDS = [0, 250, 500, 750] +EXPECTED_COT_TEMPLATE = ( + "Your task is {instruction}. To identify the key objects for your task. " + "Locate their bounding boxes in [x1,y1,x2,y2] format." +) + +PINNED_SOURCE_FILES = { + "starVLA/model/framework/VLM4A/QwenGR00T.py": + "645d99d8d6a8daaccb7bb6e3211971b5cc7396d39968b0e9c20c3894d6883249", + "starVLA/model/modules/action_model/GR00T_ActionHeader.py": + "a01c7ca048589835a23bf46cf670275dfa643a1fb2da0bafd14654e1a57236e5", + "starVLA/model/modules/action_model/flow_matching_head/cross_attention_dit.py": + "c18d2e128dddcd67dc88c4fb178c99d7ceb7ea40d40ea9622b120151b81db359", + "starVLA/model/modules/vlm/QWen3.py": + "03e0c35cfe86490886ff26a59230f27726ba4b46259d2be20beed4c532925d47", + "deployment/model_server/policy_norm_processor.py": + "3fd280c8f5072943fad6809dd5705cb713007c10d7240d2a23e3dadcd3963d2a", +} + +ARRAY_SOURCE_DTYPES = { + "input_ids": "int64", + "attention_mask": "bool", + "image_grid_thw": "int64", + "raw_l_out_35": "bfloat16", + "initial_noise": "bfloat16", + "action_features": "float32", + "dit_inputs": "float32", + "dit_block_outputs": "float32", + "dit_outputs": "float32", + "predicted_velocities": "float32", + "actions_after_steps": "float32", + "normalized_actions": "float32", + "unnormalized_actions": "float32", +} + + +def _distribution_version(name: str) -> str: + try: + return importlib.metadata.version(name) + except importlib.metadata.PackageNotFoundError as exc: + raise StarVLAError(f"required package is not installed: {name}") from exc + + +def _regular_file(path: Path, *, label: str) -> Path: + if path.is_symlink() or not path.is_file(): + raise StarVLAError(f"{label} must be a regular, non-symlink file: {path}") + return path.resolve() + + +def _source_asset_hashes(entry: Mapping[str, Any], *, staged: bool = False) -> dict[str, str]: + overrides = entry.get("staged_overrides", {}) if staged else {} + return { + relative: overrides.get(relative, record)["sha256"] + for relative, record in entry["file_hashes"].items() + } + + +def verify_source_semantics(source_dir: Path) -> dict[str, Any]: + """Bind the sampler schedule and ordering to the pinned official source.""" + + actual: dict[str, str] = {} + for relative, expected_hash in PINNED_SOURCE_FILES.items(): + path = _regular_file(source_dir / relative, label=f"pinned source {relative}") + digest = sha256_file(path) + if digest != expected_hash: + raise StarVLAError( + f"pinned GR00T source SHA256 mismatch for {relative}: " + f"expected {expected_hash}, got {digest}" + ) + actual[relative] = digest + + action_source = (source_dir / "starVLA/model/modules/action_model/GR00T_ActionHeader.py").read_text( + encoding="utf-8" + ) + qwen_source = (source_dir / "starVLA/model/framework/VLM4A/QwenGR00T.py").read_text( + encoding="utf-8" + ) + required_action_fragments = ( + "dtype=vl_embs.dtype", + "dt = 1.0 / num_steps", + "t_cont = t / float(num_steps)", + "t_discretized = int(t_cont * self.num_timestep_buckets)", + "torch.cat((future_tokens, action_features), dim=1)", + "pred_velocity = pred[:, -self.action_horizon :]", + "actions = actions + dt * pred_velocity", + ) + required_qwen_fragments = ( + "backbone_attention_mask = backbone_attention_mask.to(dtype=torch.bool)", + "last_hidden = qwenvl_outputs.hidden_states[-1]", + "self.action_model.predict_action(", + "last_hidden, state, encoder_attention_mask=backbone_attention_mask", + ) + missing = [value for value in required_action_fragments if value not in action_source] + missing += [value for value in required_qwen_fragments if value not in qwen_source] + if missing: + raise StarVLAError(f"pinned GR00T source semantics probe failed: missing {missing!r}") + return { + "files": actual, + "schedule_source": "GR00T_ActionHeader.FlowmatchingActionHead.predict_action", + "continuous_formula": "t / float(num_steps)", + "bucket_formula": "int(t_cont * self.num_timestep_buckets)", + "observed_expected_timestep_ids": EXPECTED_TIMESTEP_IDS, + "query_sequence": "future_tokens_then_action_features", + "conditioning_sequence": "complete_qwen_outer_raw_l_out_35_with_bool_attention_mask", + "velocity_slice": "last_action_horizon_tokens", + "euler_update": "actions = actions + dt * pred_velocity", + } + + +def validate_available_inputs( + *, checkpoint_root: Path, source_dir: Path, catalog_path: Path = DEFAULT_CATALOG +) -> dict[str, Any]: + catalog = load_catalog(catalog_path) + variant = get_variant(catalog, SUPPORTED_VARIANT) + qwen = catalog["shared_assets"]["qwen3_vl_4b_instruct"] + checkpoint_root = checkpoint_root.resolve() + expected_source = (checkpoint_root / "source" / "starvla").resolve() + if source_dir.resolve() != expected_source: + raise StarVLAError( + f"StarVLA source must be the canonical checkout {expected_source}, got {source_dir.resolve()}" + ) + verify_pinned_source_checkout(source_dir, catalog["source_revisions"]["starvla"]) + source_probe = verify_source_semantics(source_dir) + policy_dir = checkpoint_root / "sources" / variant["directory"] / variant["revision"] + qwen_dir = checkpoint_root / "sources" / qwen["directory"] / qwen["revision"] + verify_catalog_files(policy_dir, variant) + verify_catalog_files(qwen_dir, qwen) + checkpoint = policy_dir / variant["checkpoint"]["path"] + sidecar = Path(f"{checkpoint}.aria2") + checkpoint_ready = checkpoint.is_file() and not checkpoint.is_symlink() and not sidecar.exists() + if checkpoint_ready: + verify_checkpoint_file(checkpoint, variant) + return { + "catalog": catalog, + "catalog_path": catalog_path.resolve(), + "variant": variant, + "qwen": qwen, + "policy_dir": policy_dir, + "qwen_dir": qwen_dir, + "checkpoint": checkpoint, + "checkpoint_ready": checkpoint_ready, + "source_dir": source_dir.resolve(), + "source_probe": source_probe, + } + + +def _validate_effective_config(config: Mapping[str, Any]) -> None: + try: + framework = config["framework"] + action = framework["action_model"] + diffusion = action["diffusion_model_cfg"] + vla = config["datasets"]["vla_data"] + except (KeyError, TypeError) as exc: + raise StarVLAError("effective GR00T config is missing required objects") from exc + actual = { + "framework": framework.get("name"), + "action_model_type": action.get("action_model_type"), + "action_horizon": action.get("action_horizon"), + "action_dim": action.get("action_dim"), + "state_dim": action.get("state_dim"), + "steps": action.get("num_inference_timesteps"), + "buckets": action.get("num_timestep_buckets"), + "future_tokens": action.get("num_target_vision_tokens"), + "width": diffusion.get("input_embedding_dim"), + "layers": diffusion.get("num_layers"), + "heads": diffusion.get("num_attention_heads"), + "head_dim": diffusion.get("attention_head_dim"), + "cross_dim": diffusion.get("cross_attention_dim"), + "output_dim": diffusion.get("output_dim"), + "interleave": diffusion.get("interleave_self_attention"), + "image_size": vla.get("image_size"), + "obs_image_size": vla.get("obs_image_size"), + "obs": vla.get("obs"), + "data_mix": vla.get("data_mix"), + "cot": vla.get("CoT_prompt"), + } + expected = { + "framework": "QwenGR00T", + "action_model_type": "DiT-B", + "action_horizon": 16, + "action_dim": 7, + "state_dim": 7, + "steps": 4, + "buckets": 1000, + "future_tokens": 32, + "width": 768, + "layers": 16, + "heads": 12, + "head_dim": 64, + "cross_dim": 2560, + "output_dim": 1024, + "interleave": True, + "image_size": [224, 224], + "obs_image_size": None, + "obs": ["image_0"], + "data_mix": "bridge_rt_1", + "cot": EXPECTED_COT_TEMPLATE, + } + if actual != expected: + raise StarVLAError(f"unexpected effective official GR00T config: {actual}") + + +def expected_model_instruction(config: Mapping[str, Any], task: str) -> str: + if not isinstance(task, str) or not task or "\x00" in task: + raise StarVLAError("task must be a non-empty string without NUL") + template = config["datasets"]["vla_data"]["CoT_prompt"] + if template != EXPECTED_COT_TEMPLATE or template.count("{instruction}") != 1: + raise StarVLAError("official GR00T CoT prompt contract changed") + return template.replace("{instruction}", task) + + +def load_official_framework(paths: Mapping[str, Any], *, device: str) -> tuple[Any, dict[str, Any]]: + import torch + import transformers + + if not paths["checkpoint_ready"]: + raise StarVLAError( + f"official GR00T checkpoint is absent or incomplete: {paths['checkpoint']}" + ) + source_dir = Path(paths["source_dir"]) + if any(name == "starVLA" or name.startswith("starVLA.") for name in sys.modules): + raise StarVLAError("starVLA was imported before pinned-source verification") + sys.path.insert(0, str(source_dir)) + try: + from starVLA.model.framework import base_framework, share_tools + from starVLA.model.framework.VLM4A import QwenGR00T + + for module in (base_framework, share_tools, QwenGR00T): + _assert_module_origin(module, source_dir) + config = resolve_effective_config(Path(paths["policy_dir"]), SUPPORTED_VARIANT) + _validate_effective_config(config) + qwen_dir = Path(paths["qwen_dir"]).resolve() + with _official_qwen_model_alias(qwen_dir) as qwen_alias: + config = base_framework.merge_config_overrides( + config, + [ + f"framework.qwenvl.base_vlm={qwen_alias}", + "framework.qwenvl.attn_implementation=sdpa", + ], + ) + configured_qwen = Path(config["framework"]["qwenvl"]["base_vlm"]) + if "Qwen3-VL" not in str(configured_qwen) or configured_qwen.resolve() != qwen_dir: + raise StarVLAError( + "effective GR00T Qwen source does not preserve official dispatch and pinned assets" + ) + cfg = share_tools.dict_to_namespace(config) + cfg.trainer.pretrained_checkpoint = None + with _config_only_qwen_bootstrap(torch, transformers, qwen_dir): + framework = QwenGR00T.Qwen_GR00T(cfg) + try: + state = torch.load(paths["checkpoint"], map_location="cpu", mmap=True, weights_only=True) + except TypeError: + state = torch.load(paths["checkpoint"], map_location="cpu", weights_only=True) + if not isinstance(state, Mapping) or not state: + raise StarVLAError("official GR00T checkpoint did not contain a state_dict") + framework.load_state_dict(state, strict=True) + del state + gc.collect() + action_model = framework.action_model + if type(framework).__name__ != "Qwen_GR00T": + raise StarVLAError(f"unexpected official framework class: {type(framework).__name__}") + if len(action_model.model.transformer_blocks) != EXPECTED_DIT_BLOCK_COUNT: + raise StarVLAError("official GR00T DiT block count is not 16") + if int(framework.action_horizon) != EXPECTED_ACTION_HORIZON: + raise StarVLAError("official GR00T action horizon changed") + qwen_dtypes = {parameter.dtype for parameter in framework.qwen_vl_interface.parameters()} + policy_dtypes = {parameter.dtype for parameter in action_model.parameters()} + if qwen_dtypes != {torch.bfloat16} or policy_dtypes != {torch.float32}: + raise StarVLAError( + "official GR00T strict-load dtype boundary changed: " + f"qwen={qwen_dtypes}, policy={policy_dtypes}" + ) + return framework.to(device).eval(), config + finally: + if sys.path and sys.path[0] == str(source_dir): + del sys.path[0] + + +def _tensor_to_array(tensor: Any) -> tuple[np.ndarray, str]: + source_dtype = str(tensor.dtype).removeprefix("torch.") + value = tensor.detach().cpu().contiguous() + if source_dtype == "bfloat16": + value = value.float() + return np.ascontiguousarray(value.numpy()), source_dtype + + +def run_official_forward( + framework: Any, *, images: Sequence[Any], task: str, seed: int = SEED +) -> dict[str, Any]: + """Run pinned source with only the declared BF16-to-FP32 compatibility widens.""" + + import torch + + captures: dict[str, Any] = {} + qwen = framework.qwen_vl_interface + action_model = framework.action_model + language_model = qwen.model.model.language_model + raw_final: dict[str, Any] = {} + raw_qwen_taps: list[Any | None] = [None] * len(language_model.layers) + language_inputs: dict[str, Any] = {} + result_norm: dict[str, Any] = {} + block_outputs: list[Any] = [] + handles = [] + original_build = qwen.build_qwenvl_inputs + original_policy = action_model.predict_action + original_action_encoder = action_model.action_encoder.forward + original_dit = action_model.model.forward + original_decoder = action_model.action_decoder.forward + + def capture_build(*args: Any, **kwargs: Any): + if "qwen_inputs" in captures: + raise StarVLAError("official GR00T preprocessing ran more than once") + batch_images = kwargs.get("images", args[0] if args else None) + instructions = kwargs.get("instructions", args[1] if len(args) > 1 else None) + captures["processed_images"] = list(batch_images[0]) + captures["framework_instructions"] = list(instructions) + output = original_build(*args, **kwargs) + captures["qwen_inputs"] = { + key: value.detach() for key, value in output.items() if isinstance(value, torch.Tensor) + } + return output + + def capture_outer(_module: Any, _inputs: Any, output: Any): + hidden = getattr(output, "hidden_states", None) + if hidden is None or len(hidden) != 37: + raise StarVLAError("official Qwen outer recorder did not expose 37 hidden tuple entries") + captures["outer_raw_final"] = hidden[-1].detach().clone() + + def capture_language_inputs(_module: Any, args: Any, kwargs: Any): + if language_inputs: + raise StarVLAError("official Qwen language model ran more than once") + if args: + raise StarVLAError("official Qwen language model stopped using keyword inputs") + inputs_embeds = kwargs.get("inputs_embeds") + visual_pos_masks = kwargs.get("visual_pos_masks") + deepstack_visual_embeds = kwargs.get("deepstack_visual_embeds") + if (inputs_embeds is None or visual_pos_masks is None or + deepstack_visual_embeds is None): + raise StarVLAError("official Qwen language model omitted prepared visual inputs") + language_inputs["inputs_embeds"] = inputs_embeds.detach().clone() + language_inputs["visual_pos_masks"] = visual_pos_masks.detach().clone() + language_inputs["deepstack_visual_embeds"] = [ + value.detach().clone() for value in deepstack_visual_embeds + ] + + def capture_policy(*args: Any, **kwargs: Any): + vl_embs = args[0] if args else kwargs.get("vl_embs") + state = args[1] if len(args) > 1 else kwargs.get("state") + policy_mask = kwargs.get( + "encoder_attention_mask", args[2] if len(args) > 2 else None + ) + if state is not None: + raise StarVLAError("official GR00T oracle unexpectedly entered the state branch") + if policy_mask is None: + raise StarVLAError("official GR00T policy did not receive an attention mask") + captures["policy_qwen_input"] = vl_embs.detach().clone() + captures["policy_attention_mask"] = policy_mask.detach().clone() + original_randn = torch.randn + + def capture_randn(*randn_args: Any, **randn_kwargs: Any): + value = original_randn(*randn_args, **randn_kwargs) + if "initial_noise" in captures: + raise StarVLAError("official GR00T policy sampled initial noise more than once") + captures["initial_noise"] = value.detach().clone() + return value + + torch.randn = capture_randn + try: + output = original_policy(*args, **kwargs) + finally: + torch.randn = original_randn + captures["raw_policy"] = output.detach().clone() + return output + + def capture_action_encoder(actions: Any, timesteps: Any): + captures.setdefault("action_inputs", []).append(actions.detach().clone()) + output = original_action_encoder(actions.to(dtype=torch.float32), timesteps) + captures.setdefault("action_encoder_outputs", []).append(output.detach().clone()) + return output + + def capture_dit(*args: Any, **kwargs: Any): + hidden = kwargs.get("hidden_states", args[0] if args else None) + conditioning = kwargs.get("encoder_hidden_states", args[1] if len(args) > 1 else None) + timestep = kwargs.get("timestep", args[2] if len(args) > 2 else None) + captures.setdefault("dit_inputs", []).append(hidden.detach().clone()) + captures.setdefault("dit_conditioning_inputs", []).append(conditioning.detach().clone()) + captures.setdefault("timestep_ids", []).append(int(timestep.item())) + if "encoder_hidden_states" in kwargs: + kwargs["encoder_hidden_states"] = conditioning.to(dtype=torch.float32) + else: + args = list(args) + args[1] = conditioning.to(dtype=torch.float32) + args = tuple(args) + output = original_dit(*args, **kwargs) + captures.setdefault("dit_outputs", []).append(output.detach().clone()) + return output + + def capture_decoder(value: Any): + captures.setdefault("decoder_inputs", []).append(value.detach().clone()) + output = original_decoder(value) + captures.setdefault("decoder_outputs", []).append(output.detach().clone()) + return output + + def capture_raw_qwen_tap(layer_index: int): + def capture(_module: Any, _inputs: Any, output: Any): + value = output.detach().clone() + raw_qwen_taps[layer_index] = value + if layer_index + 1 == len(raw_qwen_taps): + raw_final["value"] = value + + return capture + + for layer_index, layer in enumerate(language_model.layers): + handles.append(layer.register_forward_hook(capture_raw_qwen_tap(layer_index))) + handles.append( + language_model.register_forward_pre_hook( + capture_language_inputs, with_kwargs=True + ) + ) + handles.append( + language_model.norm.register_forward_hook( + lambda _m, _i, output: result_norm.__setitem__("value", output.detach().clone()) + ) + ) + handles.append(qwen.model.register_forward_hook(capture_outer)) + for block in action_model.model.transformer_blocks: + handles.append( + block.register_forward_hook( + lambda _m, _i, output: block_outputs.append(output.detach().clone()) + ) + ) + qwen.build_qwenvl_inputs = capture_build + action_model.predict_action = capture_policy + action_model.action_encoder.forward = capture_action_encoder + action_model.model.forward = capture_dit + action_model.action_decoder.forward = capture_decoder + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + try: + result = framework.predict_action(examples=[{"image": list(images), "lang": task}]) + finally: + for handle in handles: + handle.remove() + qwen.build_qwenvl_inputs = original_build + action_model.predict_action = original_policy + action_model.action_encoder.forward = original_action_encoder + action_model.model.forward = original_dit + action_model.action_decoder.forward = original_decoder + + required = { + "qwen_inputs", "processed_images", "framework_instructions", "outer_raw_final", + "policy_qwen_input", "policy_attention_mask", "initial_noise", "raw_policy", "action_inputs", + "action_encoder_outputs", "dit_inputs", "dit_conditioning_inputs", "timestep_ids", + "dit_outputs", "decoder_inputs", "decoder_outputs", + } + missing = sorted(required - set(captures)) + if (missing or "value" not in raw_final or "value" not in result_norm or + any(value is None for value in raw_qwen_taps) or not language_inputs): + raise StarVLAError(f"official GR00T instrumentation missed captures: {missing}") + if captures["framework_instructions"] != [task]: + raise StarVLAError("official GR00T framework instruction changed before Qwen preprocessing") + if len(captures["processed_images"]) != len(images) or any( + actual.mode != expected.mode or actual.size != expected.size or + actual.tobytes() != expected.tobytes() + for actual, expected in zip(captures["processed_images"], images) + ): + raise StarVLAError("official GR00T unexpectedly pre-resized or altered the input image") + if captures["timestep_ids"] != EXPECTED_TIMESTEP_IDS: + raise StarVLAError( + f"official source-derived GR00T timestep order changed: {captures['timestep_ids']}" + ) + if len(block_outputs) != 4 * EXPECTED_DIT_BLOCK_COUNT: + raise StarVLAError("official GR00T instrumentation missed DiT block outputs") + for name in ( + "action_inputs", "action_encoder_outputs", "dit_inputs", "dit_conditioning_inputs", + "dit_outputs", "decoder_inputs", "decoder_outputs", + ): + if len(captures[name]) != 4: + raise StarVLAError(f"official GR00T {name} did not run exactly four times") + if not torch.equal(captures["outer_raw_final"], raw_final["value"]): + raise StarVLAError("outer hidden_states[-1] is not complete raw l_out-35") + if torch.equal(captures["outer_raw_final"], result_norm["value"]): + raise StarVLAError("GR00T conditioning unexpectedly uses result_norm") + if not torch.equal(captures["policy_qwen_input"], captures["outer_raw_final"]): + raise StarVLAError("GR00T policy did not receive complete raw l_out-35") + mask = captures["qwen_inputs"].get("attention_mask") + if mask is None or mask.dtype not in (torch.int64, torch.bool): + raise StarVLAError("official Qwen attention mask source dtype changed") + policy_mask = captures["policy_attention_mask"] + if policy_mask.dtype != torch.bool or not torch.equal(policy_mask, mask.to(dtype=torch.bool)): + raise StarVLAError("official GR00T policy mask is not the complete Qwen boolean mask") + if captures["initial_noise"].dtype != torch.bfloat16: + raise StarVLAError("official GR00T initial noise is no longer sampled as BF16") + if not torch.equal(captures["initial_noise"], captures["action_inputs"][0]): + raise StarVLAError("first GR00T action encoder input is not initial noise") + if captures["outer_raw_final"].dtype != torch.bfloat16: + raise StarVLAError("official raw l_out-35 boundary is no longer BF16") + if any(value.dtype != torch.bfloat16 for value in raw_qwen_taps): + raise StarVLAError("official raw Qwen layer taps are no longer BF16") + if not torch.equal(raw_qwen_taps[-1], captures["outer_raw_final"]): + raise StarVLAError("official raw Qwen layer taps do not end at raw l_out-35") + if any(value.dtype != torch.float32 for value in captures["action_encoder_outputs"]): + raise StarVLAError("GR00T action encoder compatibility output is not FP32") + for name in ("dit_inputs", "dit_outputs", "decoder_inputs", "decoder_outputs"): + if any(value.dtype != torch.float32 for value in captures[name]): + raise StarVLAError(f"GR00T compatibility path {name} is not FP32") + if any(value.dtype != torch.bfloat16 for value in captures["dit_conditioning_inputs"]): + raise StarVLAError("GR00T raw DiT conditioning input is not BF16 before explicit widen") + + base_embeddings = language_inputs["inputs_embeds"] + visual_pos_masks = language_inputs["visual_pos_masks"] + deepstack_visual_embeds = language_inputs["deepstack_visual_embeds"] + token_count = captures["qwen_inputs"]["input_ids"].shape[1] + if (base_embeddings.dtype != torch.bfloat16 or + tuple(base_embeddings.shape) != (1, token_count, EXPECTED_QWEN_HIDDEN_DIM) or + visual_pos_masks.dtype != torch.bool or + tuple(visual_pos_masks.shape) != (1, token_count) or + len(deepstack_visual_embeds) != 3): + raise StarVLAError("official Qwen prepared input layout changed") + visual_mask = visual_pos_masks[0] + visual_token_count = int(visual_mask.sum().item()) + prepared_embeddings = torch.zeros( + (token_count, 4, EXPECTED_QWEN_HIDDEN_DIM), + dtype=torch.bfloat16, + device=base_embeddings.device, + ) + prepared_embeddings[:, 0, :] = base_embeddings[0] + for index, value in enumerate(deepstack_visual_embeds): + if (value.dtype != torch.bfloat16 or + tuple(value.shape) != (visual_token_count, EXPECTED_QWEN_HIDDEN_DIM)): + raise StarVLAError("official Qwen DeepStack prepared input layout changed") + prepared_embeddings[visual_mask, index + 1, :] = value + if captures["action_inputs"][0].dtype != torch.bfloat16 or any( + value.dtype != torch.float32 for value in captures["action_inputs"][1:] + ): + raise StarVLAError("GR00T actions must become FP32 after the first Euler update") + + future = action_model.future_tokens.weight.unsqueeze(0) + position = action_model.position_embedding.weight[:EXPECTED_ACTION_HORIZON].unsqueeze(0) + for step in range(4): + dit_input = captures["dit_inputs"][step] + if tuple(dit_input.shape) != (1, EXPECTED_SEQUENCE_LENGTH, EXPECTED_DIT_WIDTH): + raise StarVLAError("official GR00T DiT query sequence shape changed") + if not torch.equal(dit_input[:, :EXPECTED_FUTURE_TOKEN_COUNT], future): + raise StarVLAError("official GR00T DiT query prefix is not future tokens") + expected_action_features = captures["action_encoder_outputs"][step] + position + if not torch.equal(dit_input[:, EXPECTED_FUTURE_TOKEN_COUNT:], expected_action_features): + raise StarVLAError("official GR00T DiT query suffix is not positioned action features") + if not torch.equal(captures["dit_conditioning_inputs"][step], captures["outer_raw_final"]): + raise StarVLAError("GR00T cross-attention conditioning changed across sampler steps") + if not torch.equal(captures["dit_outputs"][step], captures["decoder_inputs"][step]): + raise StarVLAError("GR00T DiT-to-velocity decoder boundary changed") + + velocities = [value[:, -EXPECTED_ACTION_HORIZON:] for value in captures["decoder_outputs"]] + actions_after = captures["action_inputs"][1:] + [captures["raw_policy"]] + previous = captures["initial_noise"].to(dtype=torch.float32) + for step, (velocity, actual) in enumerate(zip(velocities, actions_after)): + expected = previous + 0.25 * velocity + if not torch.equal(actual, expected): + raise StarVLAError(f"official GR00T Euler update mismatch at step {step}") + previous = actual + normalized = np.asarray(result.get("normalized_actions"), dtype=np.float32) + raw_policy, _ = _tensor_to_array(captures["raw_policy"]) + if normalized.shape != (1, 16, 7) or not np.array_equal(normalized, raw_policy): + raise StarVLAError("official normalized_actions differ from captured GR00T policy output") + if not np.isfinite(normalized).all(): + raise StarVLAError("official GR00T output contains non-finite values") + captures["block_outputs"] = block_outputs + captures["raw_qwen_taps"] = raw_qwen_taps + captures["prepared_embeddings"] = prepared_embeddings + captures["predicted_velocities"] = velocities + captures["actions_after_steps"] = actions_after + captures["normalized_actions"] = normalized + return captures + + +def _stack(values: Sequence[Any], *, label: str) -> tuple[np.ndarray, str]: + import torch + + if not values: + raise StarVLAError(f"cannot stack empty {label}") + dtype = values[0].dtype + if any(value.dtype != dtype for value in values): + raise StarVLAError(f"{label} has mixed source dtypes") + return _tensor_to_array(torch.stack(list(values), dim=0)) + + +def build_arrays(captures: Mapping[str, Any], unnormalized: np.ndarray) -> dict[str, np.ndarray]: + qwen_inputs = captures["qwen_inputs"] + input_ids = np.ascontiguousarray(qwen_inputs["input_ids"][0].cpu().numpy(), dtype=np.int64) + attention_mask = np.ascontiguousarray( + qwen_inputs["attention_mask"][0].to(dtype=__import__("torch").bool).cpu().numpy(), dtype=np.bool_ + ) + image_grid = np.ascontiguousarray(qwen_inputs["image_grid_thw"].cpu().numpy(), dtype=np.int64) + raw_final, _ = _tensor_to_array(captures["outer_raw_final"]) + initial_noise, _ = _tensor_to_array(captures["initial_noise"]) + dit_inputs, _ = _stack(captures["dit_inputs"], label="DiT inputs") + block_outputs, _ = _stack(captures["block_outputs"], label="DiT block outputs") + block_outputs = block_outputs.reshape( + 4, EXPECTED_DIT_BLOCK_COUNT, 1, EXPECTED_SEQUENCE_LENGTH, EXPECTED_DIT_WIDTH + ) + dit_outputs, _ = _stack(captures["dit_outputs"], label="DiT outputs") + velocities, _ = _stack(captures["predicted_velocities"], label="predicted velocities") + actions_after, _ = _stack(captures["actions_after_steps"], label="actions after steps") + action_features = np.ascontiguousarray( + dit_inputs[:, :, EXPECTED_FUTURE_TOKEN_COUNT:, :], dtype=np.float32 + ) + arrays = { + "input_ids": input_ids, + "attention_mask": attention_mask, + "image_grid_thw": image_grid, + "raw_l_out_35": np.ascontiguousarray(raw_final, dtype=np.float32), + "initial_noise": np.ascontiguousarray(initial_noise, dtype=np.float32), + "action_features": action_features, + "dit_inputs": np.ascontiguousarray(dit_inputs, dtype=np.float32), + "dit_block_outputs": np.ascontiguousarray(block_outputs, dtype=np.float32), + "dit_outputs": np.ascontiguousarray(dit_outputs, dtype=np.float32), + "predicted_velocities": np.ascontiguousarray(velocities, dtype=np.float32), + "actions_after_steps": np.ascontiguousarray(actions_after, dtype=np.float32), + "normalized_actions": np.ascontiguousarray(captures["normalized_actions"], dtype=np.float32), + "unnormalized_actions": np.ascontiguousarray(unnormalized, dtype=np.float32), + } + expected_shapes = { + "attention_mask": (input_ids.shape[0],), + "image_grid_thw": (1, 3), + "raw_l_out_35": (1, input_ids.shape[0], 2560), + "initial_noise": (1, 16, 7), + "action_features": (4, 1, 16, 768), + "dit_inputs": (4, 1, 48, 768), + "dit_block_outputs": (4, 16, 1, 48, 768), + "dit_outputs": (4, 1, 48, 1024), + "predicted_velocities": (4, 1, 16, 7), + "actions_after_steps": (4, 1, 16, 7), + "normalized_actions": (1, 16, 7), + "unnormalized_actions": (1, 16, 7), + } + for name, shape in expected_shapes.items(): + if arrays[name].shape != shape: + raise StarVLAError(f"official GR00T {name} shape mismatch: {arrays[name].shape}") + if any(not np.isfinite(value).all() for name, value in arrays.items() if name not in { + "input_ids", "attention_mask", "image_grid_thw" + }): + raise StarVLAError("official GR00T arrays contain non-finite values") + return arrays + + +def _runtime_record(torch: Any, transformers: Any, device: str, recorder: Mapping[str, Any]) -> dict[str, Any]: + cuda_device = torch.device(device) + index = cuda_device.index if cuda_device.index is not None else torch.cuda.current_device() + properties = torch.cuda.get_device_properties(index) + return { + "python": platform.python_version(), + "platform": platform.platform(), + "torch": torch.__version__, + "torchvision": _distribution_version("torchvision"), + "transformers": transformers.__version__, + "numpy": np.__version__, + "diffusers": _distribution_version("diffusers"), + "tokenizers": _distribution_version("tokenizers"), + "pillow": _distribution_version("Pillow"), + "omegaconf": _distribution_version("omegaconf"), + "accelerate": _distribution_version("accelerate"), + "safetensors": _distribution_version("safetensors"), + "official_environment_freeze": dict(OFFICIAL_ENVIRONMENT_FREEZE), + "cuda_runtime": torch.version.cuda, + "cudnn": torch.backends.cudnn.version(), + "device": str(cuda_device), + "device_name": properties.name, + "compute_capability": [properties.major, properties.minor], + "qwen3vl_recorder_probe": dict(recorder), + } + + +def _image_record(path: Path, image: Any) -> dict[str, Any]: + pixel_header = _canonical_json({"mode": image.mode, "size": list(image.size)}) + pixel_hash = hashlib.sha256(pixel_header + b"\x00" + image.tobytes()).hexdigest() + return { + "source_size": path.stat().st_size, + "source_sha256": sha256_file(path), + "decoded_mode": image.mode, + "decoded_size": list(image.size), + "decoded_pixel_sha256": pixel_hash, + } + + +def _load_images(paths: Iterable[Path]) -> tuple[list[Any], list[dict[str, Any]]]: + from PIL import Image + + images = [] + records = [] + for path in paths: + path = _regular_file(path, label="oracle image") + with Image.open(path) as opened: + opened.load() + image = opened.convert("RGB") + images.append(image) + records.append(_image_record(path, image)) + if len(images) != 1: + raise StarVLAError("official GR00T oracle requires exactly one image") + return images, records + + +def _write_runner_contract( + path: Path, *, golden_id: str, source: Mapping[str, Any], task: str, + model_instruction: str, unnorm_key: str, image_sha256: str, + array_records: Mapping[str, Any], token_count: int, initial_noise: np.ndarray, +) -> str: + raw_noise = np.ascontiguousarray(initial_noise, dtype=" Path: + import torch + import transformers + + output_dir = output_dir.resolve() + if output_dir.exists(): + raise StarVLAError(f"golden output directory already exists: {output_dir}") + output_dir.parent.mkdir(parents=True, exist_ok=True) + arrays = build_arrays(captures, unnormalized) + records = { + name: _array_record(value, source_dtype=ARRAY_SOURCE_DTYPES[name]) + for name, value in arrays.items() + } + model_instruction = expected_model_instruction(config, task) + identity = { + "schema_version": 1, + "variant": SUPPORTED_VARIANT, + "checkpoint_sha256": paths["variant"]["checkpoint"]["sha256"], + "starvla_revision": paths["catalog"]["source_revisions"]["starvla"], + "qwen_revision": paths["qwen"]["revision"], + "task": task, + "unnorm_key": unnorm_key, + "seed": SEED, + "images": [record["source_sha256"] for record in image_records], + } + golden_id = _sha256_bytes(_canonical_json(identity)) + variant = paths["variant"] + qwen = paths["qwen"] + source = { + "catalog_sha256": sha256_file(paths["catalog_path"]), + "bundle_uuid": official_bundle_uuid(variant, paths["catalog"]), + "starvla_repo_revision": paths["catalog"]["source_revisions"]["starvla"], + "checkpoint_repo_id": variant["repo_id"], + "checkpoint_revision": variant["revision"], + "checkpoint_size": variant["checkpoint"]["size"], + "checkpoint_sha256": variant["checkpoint"]["sha256"], + "policy_assets": _source_asset_hashes(variant), + "qwen_repo_id": qwen["repo_id"], + "qwen_revision": qwen["revision"], + "qwen_runtime_assets": _source_asset_hashes(qwen), + "qwen_converted_component_assets": _source_asset_hashes(qwen, staged=True), + "pinned_source_probe": paths["source_probe"], + } + with tempfile.TemporaryDirectory(prefix=f".{output_dir.name}.", dir=output_dir.parent) as temporary: + staging = Path(temporary) + inputs = staging / "inputs" + inputs.mkdir() + image_artifacts = [] + for index, (source_path, record) in enumerate(zip(image_paths, image_records)): + suffix = source_path.suffix.lower() or ".img" + destination = inputs / f"image-{index:02d}{suffix}" + shutil.copyfile(source_path, destination) + image_artifacts.append({ + **record, + "artifact": destination.relative_to(staging).as_posix(), + "artifact_size": destination.stat().st_size, + "artifact_sha256": sha256_file(destination), + }) + tensor_path = staging / "tensors.npz" + np.savez(tensor_path, **arrays) + contract_path = staging / "runner_contract.txt" + contract_sha = _write_runner_contract( + contract_path, golden_id=golden_id, source=source, task=task, + model_instruction=model_instruction, unnorm_key=unnorm_key, + image_sha256=image_records[0]["source_sha256"], array_records=records, + token_count=int(arrays["input_ids"].shape[0]), initial_noise=arrays["initial_noise"], + ) + noise_path = staging / "initial_noise.f32" + manifest: dict[str, Any] = { + "schema_version": GOLDEN_SCHEMA_VERSION, + "kind": GOLDEN_KIND, + "golden_id": golden_id, + "created_utc": dt.datetime.now(dt.timezone.utc).isoformat(), + "variant": SUPPORTED_VARIANT, + "model_type": "starvla", + "source": source, + "runtime": _runtime_record(torch, transformers, str(next(framework.parameters()).device), recorder_probe), + "determinism": { + "seed": SEED, + "rng_reset_immediately_before_predict": True, + "initial_noise_saved_explicitly": True, + "cross_language_seed_replay_allowed": False, + "torch_deterministic_algorithms": True, + "cublas_workspace_config": ":4096:8", + "cuda_matmul_allow_tf32": False, + "cudnn_allow_tf32": False, + "cudnn_benchmark": False, + "attention_implementation": "sdpa", + }, + "compatibility": { + "qwen_bootstrap": "config_only_then_strict_official_checkpoint_load", + "policy_parameters_after_strict_load": "float32", + "raw_qwen_conditioning": "bfloat16_complete_l_out_35", + "initial_noise": "torch_randn_bfloat16_then_exact_widen_at_action_encoder", + "dit_conditioning": "bfloat16_raw_l_out_35_exact_widen_to_float32", + "actions_dtype_by_step_input": ["bfloat16", "float32", "float32", "float32"], + "reason": ( + "pinned source requests CUDA autocast(dtype=float32), disabled by PyTorch 2.6; " + "the two explicit widens realize the declared FP32 action-policy path" + ), + }, + "input": { + "task": task, + "unnorm_key": unnorm_key, + "state": None, + "images": image_artifacts, + }, + "prompt": { + "framework_instruction": task, + "model_instruction": model_instruction, + "action_token_mode": "none", + }, + "model_contract": { + "action_horizon": 16, + "action_dim": 7, + "qwen_hidden_dim": 2560, + "qwen_tap": "outer_hidden_states_last_equals_complete_raw_l_out_35", + "attention_mask": "full_sequence_bool_nonzero_participates", + "future_token_count": 32, + "query_sequence_length": 48, + "query_token_order": "future_tokens_then_action_tokens", + "dit_width": 768, + "dit_output_dim": 1024, + "dit_block_count": 16, + "cross_attention_blocks": list(range(0, 16, 2)), + "self_attention_blocks": list(range(1, 16, 2)), + "timestep_ids": EXPECTED_TIMESTEP_IDS, + "euler_dt": 0.25, + "state_input_active": False, + "tap_layouts": { + "raw_l_out_35": "batch_token_hidden", + "action_features": "step_batch_action_width", + "dit_inputs": "step_batch_query_width", + "dit_block_outputs": "step_block_batch_query_width", + "dit_outputs": "step_batch_query_output", + "predicted_velocities": "step_batch_action_dimension", + "actions_after_steps": "step_batch_action_dimension", + }, + }, + "tokens": { + "input_ids": arrays["input_ids"].tolist(), + "attention_mask": arrays["attention_mask"].astype(np.uint8).tolist(), + "image_grid_thw": arrays["image_grid_thw"].tolist(), + }, + "outputs": { + "normalized_actions": arrays["normalized_actions"].tolist(), + "unnormalized_actions": arrays["unnormalized_actions"].tolist(), + }, + "artifacts": { + "tensors": { + "path": tensor_path.name, + "size": tensor_path.stat().st_size, + "sha256": sha256_file(tensor_path), + "encoding": "numpy_npz_stored", + "arrays": records, + }, + "initial_noise_raw": { + "path": noise_path.name, + "size": noise_path.stat().st_size, + "sha256": sha256_file(noise_path), + "encoding": "little_endian_float32_exact_widened_bfloat16", + "array_sha256": records["initial_noise"]["sha256"], + }, + "runner_contract": { + "path": contract_path.name, + "size": contract_path.stat().st_size, + "sha256": contract_sha, + "encoding": "ordered_utf8_key_value_v1", + }, + }, + } + manifest["integrity"] = { + "canonicalization": "utf8_json_sort_keys_compact_excluding_integrity", + "manifest_payload_sha256": _sha256_bytes(_canonical_json(manifest)), + } + (staging / "golden.json").write_text( + json.dumps(manifest, ensure_ascii=False, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + Path(temporary).replace(output_dir) + return output_dir / "golden.json" + + +def _preflight_record(paths: Mapping[str, Any], recorder: Mapping[str, Any]) -> dict[str, Any]: + config = resolve_effective_config(Path(paths["policy_dir"]), SUPPORTED_VARIANT) + _validate_effective_config(config) + return { + "variant": SUPPORTED_VARIANT, + "checkpoint": str(paths["checkpoint"]), + "checkpoint_ready": paths["checkpoint_ready"], + "expected_checkpoint_size": paths["variant"]["checkpoint"]["size"], + "expected_checkpoint_sha256": paths["variant"]["checkpoint"]["sha256"], + "source_probe": paths["source_probe"], + "qwen3vl_recorder_probe": recorder, + "effective_config_valid": True, + } + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--checkpoint-root", type=Path, default=Path("ckpts/starvla")) + parser.add_argument("--starvla-source", type=Path, default=Path("ckpts/starvla/source/starvla")) + parser.add_argument("--catalog", type=Path, default=DEFAULT_CATALOG) + parser.add_argument("--image", type=Path, action="append", default=[]) + parser.add_argument("--task", default="grab the block.") + parser.add_argument("--unnorm-key", default="oxe_bridge") + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--output-dir", type=Path, default=Path("goldens/starvla/groot/bridge-grab-block")) + parser.add_argument( + "--qwen-layer-diagnostic", + type=Path, + help="optionally write 36 x token_count x 2560 raw decoder taps as little-endian FP32", + ) + parser.add_argument( + "--qwen-prepared-embeddings", + type=Path, + help=("optionally write token_count x 10240 prepared decoder inputs as " + "little-endian FP32 exact-widened BF16"), + ) + parser.add_argument("--preflight-only", action="store_true") + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + args = build_parser().parse_args(argv) + try: + if not sys.flags.isolated: + raise StarVLAError( + "the GR00T oracle must run in isolated mode; invoke with `python -I`" + ) + import torch + import transformers + + validate_runtime_versions( + torch_version=torch.__version__, + torchvision_version=_distribution_version("torchvision"), + transformers_version=transformers.__version__, + numpy_version=np.__version__, + diffusers_version=_distribution_version("diffusers"), + tokenizers_version=_distribution_version("tokenizers"), + pillow_version=_distribution_version("Pillow"), + omegaconf_version=_distribution_version("omegaconf"), + accelerate_version=_distribution_version("accelerate"), + safetensors_version=_distribution_version("safetensors"), + ) + _configure_determinism(torch, seed=SEED, device=args.device) + paths = validate_available_inputs( + checkpoint_root=args.checkpoint_root, + source_dir=args.starvla_source, + catalog_path=args.catalog, + ) + recorder = verify_transformers_qwen3vl_recorder_semantics(torch, transformers) + if args.preflight_only: + print(json.dumps(_preflight_record(paths, recorder), indent=2, sort_keys=True)) + return 0 + if not paths["checkpoint_ready"]: + raise StarVLAError( + f"official GR00T checkpoint is not ready: {paths['checkpoint']}" + ) + if len(args.image) != 1: + raise StarVLAError("exactly one --image is required") + images, image_records = _load_images(args.image) + framework, config = load_official_framework(paths, device=args.device) + captures = run_official_forward(framework, images=images, task=args.task) + if args.qwen_layer_diagnostic is not None: + diagnostic_path = args.qwen_layer_diagnostic.resolve() + if not diagnostic_path.parent.is_dir() or diagnostic_path.exists(): + raise StarVLAError( + "Qwen layer diagnostic parent must exist and output must be absent: " + f"{diagnostic_path}" + ) + raw_qwen_taps, source_dtype = _stack( + captures["raw_qwen_taps"], label="raw Qwen layer taps" + ) + token_count = captures["qwen_inputs"]["input_ids"].shape[1] + expected_source_shape = (36, 1, token_count, 2560) + if (source_dtype != "bfloat16" or + raw_qwen_taps.shape != expected_source_shape): + raise StarVLAError( + "official raw Qwen layer diagnostic has an incompatible dtype or shape" + ) + raw_qwen_taps = raw_qwen_taps[:, 0] + with diagnostic_path.open("xb") as stream: + stream.write(np.ascontiguousarray(raw_qwen_taps, dtype=" None: + if not sys.flags.isolated: + raise StarVLAError( + "the GR00T reference server must run in isolated mode; invoke with `python -I`" + ) + + +def _validate_reference_runtime(torch: Any, transformers: Any) -> None: + validate_runtime_versions( + torch_version=torch.__version__, + torchvision_version=_distribution_version("torchvision"), + transformers_version=transformers.__version__, + numpy_version=np.__version__, + diffusers_version=_distribution_version("diffusers"), + tokenizers_version=_distribution_version("tokenizers"), + pillow_version=_distribution_version("Pillow"), + omegaconf_version=_distribution_version("omegaconf"), + accelerate_version=_distribution_version("accelerate"), + safetensors_version=_distribution_version("safetensors"), + ) + + +def install_groot_dtype_bridge(framework: Any, torch: Any) -> None: + """Install the two explicit widens used by the independent GR00T oracle.""" + + if getattr(framework, "_robotcpp_groot_dtype_bridge", False): + return + action_model = framework.action_model + policy_dtype = next(action_model.parameters()).dtype + if policy_dtype != torch.float32: + raise StarVLAError(f"official GR00T policy must remain float32, got {policy_dtype}") + + original_action_encoder = action_model.action_encoder.forward + original_dit = action_model.model.forward + + def action_encoder_with_policy_dtype(actions: Any, timesteps: Any): + if actions.dtype not in (torch.bfloat16, torch.float32): + raise StarVLAError(f"unexpected GR00T action dtype: {actions.dtype}") + return original_action_encoder(actions.to(dtype=torch.float32), timesteps) + + def dit_with_policy_dtype(*args: Any, **kwargs: Any): + conditioning = kwargs.get( + "encoder_hidden_states", args[1] if len(args) > 1 else None + ) + if conditioning is None: + raise StarVLAError("official GR00T DiT omitted encoder_hidden_states") + if conditioning.dtype != torch.bfloat16: + raise StarVLAError( + f"official GR00T Qwen conditioning must be bfloat16, got {conditioning.dtype}" + ) + if "encoder_hidden_states" in kwargs: + kwargs["encoder_hidden_states"] = conditioning.to(dtype=torch.float32) + else: + mutable_args = list(args) + mutable_args[1] = conditioning.to(dtype=torch.float32) + args = tuple(mutable_args) + return original_dit(*args, **kwargs) + + action_model.action_encoder.forward = action_encoder_with_policy_dtype + action_model.model.forward = dit_with_policy_dtype + framework._robotcpp_groot_dtype_bridge = True + + +def build_preflight_record(paths: Mapping[str, Any]) -> dict[str, Any]: + variant = paths["variant"] + qwen = paths["qwen"] + return { + "schema_version": 1, + "ready": bool(paths["checkpoint_ready"]), + "variant": variant["_catalog_key"], + "model_type": variant["model_type"], + "framework": variant["framework"], + "backbone": variant.get("backbone", "qwen3_vl"), + "checkpoint": { + "repo_id": variant["repo_id"], + "revision": variant["revision"], + "path": str(Path(paths["checkpoint"]).resolve()), + "size": int(variant["checkpoint"]["size"]), + "sha256": variant["checkpoint"]["sha256"], + }, + "qwen": { + "repo_id": qwen["repo_id"], + "revision": qwen["revision"], + "path": str(Path(paths["qwen_dir"]).resolve()), + }, + "starvla": { + "revision": paths["catalog"]["source_revisions"]["starvla"], + "path": str(Path(paths["source_dir"]).resolve()), + }, + "catalog": { + "path": str(Path(paths["catalog_path"]).resolve()), + "sha256": sha256_file(Path(paths["catalog_path"])), + }, + } + + +def build_server_metadata( + paths: Mapping[str, Any], + framework: Any, + *, + default_unnorm_key: str, + source_tree_sha1: str, + source_tracked_index_sha256: str, + runtime: Mapping[str, Any], +) -> dict[str, Any]: + catalog = paths["catalog"] + variant = paths["variant"] + qwen = paths["qwen"] + statistics_path = Path(paths["policy_dir"]) / "dataset_statistics.json" + try: + statistics = json.loads(statistics_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise StarVLAError( + f"failed to read official GR00T normalization profiles: {exc}" + ) from exc + if not isinstance(statistics, Mapping): + raise StarVLAError("official GR00T dataset_statistics.json must be an object") + profiles = [str(value) for value in statistics.keys()] + if profiles != EXPECTED_PROFILES: + raise StarVLAError( + f"unexpected GR00T normalization profiles: expected {EXPECTED_PROFILES}, got {profiles}" + ) + if default_unnorm_key not in profiles: + raise StarVLAError( + f"default unnorm key {default_unnorm_key!r} is not in {profiles}" + ) + + qwen_dtypes = sorted( + { + str(parameter.dtype).removeprefix("torch.") + for parameter in framework.qwen_vl_interface.parameters() + } + ) + policy_dtypes = sorted( + { + str(parameter.dtype).removeprefix("torch.") + for parameter in framework.action_model.parameters() + } + ) + if qwen_dtypes != ["bfloat16"] or policy_dtypes != ["float32"]: + raise StarVLAError( + f"unexpected loaded dtype profile: qwen={qwen_dtypes}, groot={policy_dtypes}" + ) + + chunk_size = int(framework.action_horizon) + action_dim = int(framework.action_model.action_dim) + if (chunk_size, action_dim) != (16, 7): + raise StarVLAError( + f"unexpected official GR00T action contract: chunk={chunk_size}, dim={action_dim}" + ) + + checkpoint_sha256 = str(variant["checkpoint"]["sha256"]) + checkpoint_revision = str(variant["revision"]) + starvla_revision = str(catalog["source_revisions"]["starvla"]) + qwen_manifest = _asset_manifest(qwen) + policy_manifest = _asset_manifest(variant) + model_info = { + "model_type": MODEL_TYPE, + "framework": FRAMEWORK, + "bundle_uuid": official_bundle_uuid(variant, catalog), + "checkpoint_sha256": checkpoint_sha256, + "checkpoint_revision": checkpoint_revision, + "qwen_revision": str(qwen["revision"]), + "starvla_revision": starvla_revision, + "image_names": [DEFAULT_IMAGE_NAME], + "state_supported": False, + "state_dimension_dynamic": False, + "state_dim": 0, + "chunk_size": chunk_size, + "action_dim": action_dim, + "normalization_profiles": profiles, + "default_unnorm_key": default_unnorm_key, + } + return { + "schema_version": SERVER_METADATA_SCHEMA_VERSION, + "protocol_version": wire.VERSION, + "backend": REFERENCE_BACKEND, + "purpose": REFERENCE_PURPOSE, + "catalog_variant": SUPPORTED_VARIANT, + "backbone": "qwen3_vl", + "runtime": dict(runtime), + "model_info": model_info, + "checkpoint": { + "repo_id": variant["repo_id"], + "revision": checkpoint_revision, + "path": str(Path(paths["checkpoint"]).resolve()), + "size": int(variant["checkpoint"]["size"]), + "sha256": checkpoint_sha256, + "asset_manifest_sha256": _canonical_sha256(policy_manifest), + }, + "qwen": { + "repo_id": qwen["repo_id"], + "revision": qwen["revision"], + "bootstrap_assets_manifest_sha256": _canonical_sha256(qwen_manifest), + "bootstrap_assets": qwen_manifest["files"], + }, + "starvla_source": { + "revision": starvla_revision, + "commit_sha": starvla_revision, + "git_tree_sha1": source_tree_sha1, + "tracked_index_manifest_sha256": source_tracked_index_sha256, + "path": str(Path(paths["source_dir"]).resolve()), + }, + "catalog": { + "path": str(Path(paths["catalog_path"]).resolve()), + "sha256": sha256_file(Path(paths["catalog_path"])), + }, + "dtype_profile": { + "qwen_parameters": "bfloat16", + "qwen_conditioning": "bfloat16", + "action_noise_initial": "bfloat16", + "action_encoder_input_cast": "float32", + "dit_conditioning_cast": "float32", + "groot_parameters": "float32", + "wire_actions": "float32", + "whole_model_cast": False, + }, + "action_contract": {"chunk_size": chunk_size, "action_dim": action_dim}, + "normalization": { + "implementation": "official PolicyNormProcessor", + "available_unnorm_keys": profiles, + "default_unnorm_key": default_unnorm_key, + "runtime_robot_profile_aliases": {profile: profile for profile in profiles}, + }, + "runtime_version_contract": { + "torch": EXPECTED_TORCH_VERSION, + "torchvision": EXPECTED_TORCHVISION_VERSION, + "transformers": EXPECTED_TRANSFORMERS_VERSION, + "numpy": EXPECTED_NUMPY_VERSION, + "diffusers": EXPECTED_DIFFUSERS_VERSION, + "tokenizers": EXPECTED_TOKENIZERS_VERSION, + "pillow": EXPECTED_PILLOW_VERSION, + "omegaconf": EXPECTED_OMEGACONF_VERSION, + "accelerate": EXPECTED_ACCELERATE_VERSION, + "safetensors": EXPECTED_SAFETENSORS_VERSION, + }, + } + + +class PinnedGROOTReferencePolicy: + """Original-checkpoint GR00T inference plus the official unnormalizer.""" + + def __init__( + self, + *, + framework: Any, + processor_factory: Callable[..., Any], + checkpoint: Path, + metadata: Mapping[str, Any], + ) -> None: + self.framework = framework + self.metadata = dict(metadata) + self.model_info = dict(self.metadata["model_info"]) + self.unnorm_key = str(self.model_info["default_unnorm_key"]) + self.processor = processor_factory( + str(checkpoint), unnorm_key=self.unnorm_key + ) + if self.processor.unnorm_key != self.unnorm_key: + raise StarVLAError("PolicyNormProcessor selected the wrong profile") + + def reset(self) -> None: + # GR00T has no observation history. Preserve the seeded noise stream. + return None + + def predict(self, request: PredictRequest) -> PredictResult: + if len(request.images) != 1: + raise ProtocolError( + f"GR00T Bridge reference requires exactly one image, got {len(request.images)}" + ) + image = request.images[0] + if image.name != DEFAULT_IMAGE_NAME: + raise ProtocolError( + f"GR00T Bridge reference requires image name {DEFAULT_IMAGE_NAME!r}, got {image.name!r}" + ) + if request.state: + raise ProtocolError("GR00T Bridge reference does not accept robot state") + if not request.task.strip(): + raise ProtocolError("task must not be empty") + + try: + from PIL import Image + except ImportError as exc: + raise RuntimeError("Pillow is required for GR00T reference inference") from exc + + pil_image = Image.fromarray(image.to_rgb_array(), mode="RGB") + total_started = time.perf_counter() + forward_started = time.perf_counter() + output = self.framework.predict_action( + examples=[{"image": [pil_image], "lang": request.task}] + ) + forward_ms = (time.perf_counter() - forward_started) * 1000.0 + if not isinstance(output, Mapping) or "normalized_actions" not in output: + raise RuntimeError("official GR00T forward did not return normalized_actions") + normalized = np.asarray(output["normalized_actions"]) + expected_shape = ( + 1, + int(self.model_info["chunk_size"]), + int(self.model_info["action_dim"]), + ) + if normalized.shape != expected_shape or not np.isfinite(normalized).all(): + raise RuntimeError( + f"official GR00T returned invalid normalized actions: {normalized.shape}" + ) + + unnorm_started = time.perf_counter() + actions = np.asarray( + self.processor.unapply_actions(normalized[0]), + dtype=np.float32, + ) + unnorm_ms = (time.perf_counter() - unnorm_started) * 1000.0 + if actions.shape != expected_shape[1:] or not np.isfinite(actions).all(): + raise RuntimeError( + f"official PolicyNormProcessor returned invalid actions: {actions.shape}" + ) + return PredictResult( + actions=np.ascontiguousarray(actions), + metrics={ + "python_forward_ms": forward_ms, + "python_unnorm_ms": unnorm_ms, + "model_total_ms": (time.perf_counter() - total_started) * 1000.0, + }, + ) + + +def load_pinned_reference_policy( + *, + checkpoint_root: Path, + starvla_source: Path | None, + device: str, + noise_seed: int, + default_unnorm_key: str, +) -> PinnedGROOTReferencePolicy: + source_dir = starvla_source or checkpoint_root / "source" / "starvla" + paths = validate_available_inputs( + checkpoint_root=checkpoint_root, + source_dir=Path(source_dir), + catalog_path=DEFAULT_CATALOG, + ) + try: + import torch + import transformers + except ImportError as exc: + raise StarVLAError(f"official StarVLA runtime dependency is missing: {exc}") from exc + _validate_reference_runtime(torch, transformers) + _configure_determinism(torch, seed=noise_seed, device=device) + framework, _config = load_official_framework(paths, device=device) + install_groot_dtype_bridge(framework, torch) + + source_dir = Path(paths["source_dir"]) + sys.path.insert(0, str(source_dir)) + try: + from deployment.model_server import policy_norm_processor + + _assert_module_origin(policy_norm_processor, source_dir) + processor_factory = policy_norm_processor.PolicyNormProcessor + finally: + if sys.path and sys.path[0] == str(source_dir): + del sys.path[0] + + metadata = build_server_metadata( + paths, + framework, + default_unnorm_key=default_unnorm_key, + source_tree_sha1=_git_tree_sha1(source_dir), + source_tracked_index_sha256=_git_tracked_index_sha256(source_dir), + runtime=build_runtime_metadata(torch, transformers, device=device), + ) + return PinnedGROOTReferencePolicy( + framework=framework, + processor_factory=processor_factory, + checkpoint=Path(paths["checkpoint"]), + metadata=metadata, + ) + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--variant", choices=(SUPPORTED_VARIANT,), default=SUPPORTED_VARIANT) + parser.add_argument("--checkpoint-root", type=Path, default=Path("ckpts/starvla")) + parser.add_argument("--starvla-source", type=Path) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, default=5555) + parser.add_argument("--unnorm-key", default=DEFAULT_UNNORM_KEY) + parser.add_argument("--noise-seed", type=int, default=0) + parser.add_argument("--metadata-output", type=Path) + parser.add_argument("--preflight", action="store_true") + parser.add_argument("--verbosity", type=int, default=0) + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + args = build_parser().parse_args(argv) + _require_isolated_python() + if args.host != "127.0.0.1": + raise StarVLAError("--host must be 127.0.0.1") + if args.port <= 0 or args.port > 65535: + raise StarVLAError("--port must be in 1..65535") + if args.noise_seed < 0: + raise StarVLAError("--noise-seed must be non-negative") + if args.verbosity < 0: + raise StarVLAError("--verbosity must be non-negative") + logging.basicConfig( + level=logging.DEBUG if args.verbosity else logging.INFO, + format="%(asctime)s %(levelname)s %(message)s", + force=True, + ) + + checkpoint_root = args.checkpoint_root.resolve() + source_dir = ( + args.starvla_source.resolve() + if args.starvla_source + else checkpoint_root / "source" / "starvla" + ) + if args.preflight: + paths = validate_available_inputs( + checkpoint_root=checkpoint_root, + source_dir=source_dir, + catalog_path=DEFAULT_CATALOG, + ) + record = build_preflight_record(paths) + if not record["ready"]: + raise StarVLAError(f"official GR00T checkpoint is incomplete: {paths['checkpoint']}") + if args.metadata_output is not None: + write_metadata(args.metadata_output, record) + sys.stdout.write(json.dumps(record, indent=2, sort_keys=True) + "\n") + return 0 + + policy = load_pinned_reference_policy( + checkpoint_root=checkpoint_root, + starvla_source=source_dir, + device=args.device, + noise_seed=args.noise_seed, + default_unnorm_key=args.unnorm_key, + ) + if args.metadata_output is not None: + write_metadata(args.metadata_output, policy.metadata) + logging.info( + "loaded pinned GR00T Python reference metadata=%s", + _canonical_json_bytes(policy.metadata).decode("ascii"), + ) + server = ReferenceProtocolServer(policy, host=args.host, port=args.port) + logging.info( + "Python reference server listening on %s:%d model=%s variant=%s", + server.address[0], + server.address[1], + MODEL_TYPE, + args.variant, + ) + try: + server.serve_forever() + except KeyboardInterrupt: + logging.info("Python reference server interrupted") + server.close() + return 0 + + +if __name__ == "__main__": + try: + raise SystemExit(main()) + except StarVLAError as exc: + raise SystemExit(f"error: {exc}") from exc From 695488b280b04d945006f99537515eb007ad1b4c Mon Sep 17 00:00:00 2001 From: JJJYmmm <1650675829@qq.com> Date: Mon, 10 Aug 2026 12:42:56 +0800 Subject: [PATCH 07/11] starvla: add Qwen2.5-VL GR00T policy --- .../generate_starvla_qwen25_groot_golden.py | 1142 +++++++++++++++++ 1 file changed, 1142 insertions(+) create mode 100644 tools/hf2gguf/starvla/generate_starvla_qwen25_groot_golden.py diff --git a/tools/hf2gguf/starvla/generate_starvla_qwen25_groot_golden.py b/tools/hf2gguf/starvla/generate_starvla_qwen25_groot_golden.py new file mode 100644 index 0000000..0c6c461 --- /dev/null +++ b/tools/hf2gguf/starvla/generate_starvla_qwen25_groot_golden.py @@ -0,0 +1,1142 @@ +#!/usr/bin/env python3 +"""Generate a fixed-noise local-Python oracle for official Qwen2.5-VL GR00T. + +The oracle is intentionally separate from the Qwen3-VL GR00T schema. The +released Qwen2.5 model conditions its flow head on hidden tuple entry 36, +which is the final ``result_norm`` tensor, while the Qwen3 model uses the raw +outer ``l_out-35`` recorder boundary. +""" + +from __future__ import annotations + +import argparse +import datetime as dt +import gc +import json +import os +import random +import shutil +import sys +import tempfile +from pathlib import Path +from typing import Any, Iterable, Mapping, Sequence + +import numpy as np + + +TOOLS_DIR = Path(__file__).resolve().parent +if str(TOOLS_DIR) not in sys.path: + sys.path.insert(0, str(TOOLS_DIR)) + +from generate_starvla_oft_golden import ( # noqa: E402 + _array_record, + _assert_module_origin, + _canonical_json, + _configure_determinism, + _distribution_version, + _ensure_regular_file, + _image_pixel_sha256, + _require_isolated_python, + _runtime_record, + _sha256_bytes, + _tensor_to_array, + validate_runtime_versions, +) +from generate_starvla_qwen25_oft_golden import ( # noqa: E402 + EXPECTED_QWEN_VL_UTILS_VERSION, + _config_only_qwen25_bootstrap, + _official_qwen25_alias, + _qwen_asset_records, + _verify_clean_source, +) +from starvla_checkpoint import ( # noqa: E402 + DEFAULT_CATALOG, + StarVLAError, + get_variant, + load_catalog, + official_bundle_uuid, + sha256_file, + verify_catalog_files, + verify_checkpoint_file, +) + + +SCHEMA_VERSION = 1 +GOLDEN_KIND = "starvla_qwen25_groot_local_pt_python_oracle" +MODEL_TYPE = "starvla" +VARIANT = "qwen25_groot" +BACKBONE = "qwen2_5_vl" +ACTION_RELATIVE_L2_LIMIT = 0.03 +SEED = 0 + +OFFICIAL_CHECKPOINT_REPO_ID = "StarVLA/Qwen-GR00T-Bridge-RT-1" +OFFICIAL_CHECKPOINT_REVISION = "5ebc661ba38b29c28f20fff6574801e6f49f3466" +OFFICIAL_CHECKPOINT_FILENAME = "steps_30000_pytorch_model.pt" +OFFICIAL_CHECKPOINT_SIZE = 8_456_891_339 +OFFICIAL_CHECKPOINT_SHA256 = "9646da2ae0b32589a75c8cc88fae96c93c5d269b69fd7a29200744936e01d96f" +OFFICIAL_QWEN_REPO_ID = "StarVLA/Qwen2.5-VL-3B-Instruct-Action" +OFFICIAL_QWEN_REVISION = "ce86bd9a53416527b8361e8dfc47316288ffa110" +OFFICIAL_STARVLA_REPO_ID = "starVLA/starVLA" +OFFICIAL_STARVLA_REVISION = "631aae02afe6d95876e923ff518e8ff2ab9a2f88" + +EXPECTED_ACTION_HORIZON = 16 +EXPECTED_ACTION_DIM = 7 +EXPECTED_QWEN_HIDDEN_DIM = 2048 +EXPECTED_QWEN_LAYER_COUNT = 36 +EXPECTED_HIDDEN_TUPLE_INDEX = 36 +EXPECTED_DIT_WIDTH = 768 +EXPECTED_DIT_OUTPUT_DIM = 1024 +EXPECTED_DIT_BLOCK_COUNT = 16 +EXPECTED_FUTURE_TOKEN_COUNT = 32 +EXPECTED_TIMESTEP_IDS = [0, 250, 500, 750] +EXPECTED_COT_TEMPLATE = ( + "Your task is {instruction}. To identify the key objects for your task. " + "Locate their bounding boxes in [x1,y1,x2,y2] format." +) +ACTION_TOKEN_ID_MIN = 151665 +ACTION_TOKEN_ID_MAX = 153712 +ACTION_TOKEN_COUNT = 2048 +UNNORM_KEYS = ("oxe_bridge", "oxe_rt1") + +PINNED_SOURCE_FILES = { + "starVLA/model/framework/VLM4A/QwenGR00T.py": + "645d99d8d6a8daaccb7bb6e3211971b5cc7396d39968b0e9c20c3894d6883249", + "starVLA/model/modules/action_model/GR00T_ActionHeader.py": + "a01c7ca048589835a23bf46cf670275dfa643a1fb2da0bafd14654e1a57236e5", + "starVLA/model/modules/action_model/flow_matching_head/cross_attention_dit.py": + "c18d2e128dddcd67dc88c4fb178c99d7ceb7ea40d40ea9622b120151b81db359", + "starVLA/model/modules/vlm/QWen2_5.py": + "296a6b22859517ed9c302bc7c4e1c3362690e1da12ec8ad9a26b8a90d25dabec", + "deployment/model_server/policy_norm_processor.py": + "3fd280c8f5072943fad6809dd5705cb713007c10d7240d2a23e3dadcd3963d2a", +} + + +def _load_json_object(path: Path, *, label: str) -> dict[str, Any]: + _ensure_regular_file(path, label=label) + try: + value = json.loads(path.read_text(encoding="utf-8")) + except (OSError, UnicodeError, json.JSONDecodeError) as exc: + raise StarVLAError(f"failed to parse {label} {path}: {exc}") from exc + if not isinstance(value, dict): + raise StarVLAError(f"{label} root must be an object") + return value + + +def verify_source_semantics(source_dir: Path) -> dict[str, Any]: + actual: dict[str, str] = {} + for relative, expected in PINNED_SOURCE_FILES.items(): + path = source_dir / relative + _ensure_regular_file(path, label=f"pinned source {relative}") + digest = sha256_file(path) + if digest != expected: + raise StarVLAError( + f"pinned Qwen2.5 GR00T source SHA256 mismatch for {relative}: " + f"expected {expected}, got {digest}" + ) + actual[relative] = digest + + framework_source = ( + source_dir / "starVLA/model/framework/VLM4A/QwenGR00T.py" + ).read_text(encoding="utf-8") + action_source = ( + source_dir / "starVLA/model/modules/action_model/GR00T_ActionHeader.py" + ).read_text(encoding="utf-8") + required = ( + "last_hidden = qwenvl_outputs.hidden_states[-1]", + "backbone_attention_mask = backbone_attention_mask.to(dtype=torch.bool)", + "last_hidden, state, encoder_attention_mask=backbone_attention_mask", + "dtype=vl_embs.dtype", + "t_cont = t / float(num_steps)", + "t_discretized = int(t_cont * self.num_timestep_buckets)", + "actions = actions + dt * pred_velocity", + ) + combined = framework_source + "\n" + action_source + missing = [fragment for fragment in required if fragment not in combined] + if missing: + raise StarVLAError( + f"pinned Qwen2.5 GR00T source semantics probe failed: {missing!r}" + ) + return { + "files": actual, + "hidden_selection": "qwenvl_outputs.hidden_states[-1]", + "hidden_tuple_index": EXPECTED_HIDDEN_TUPLE_INDEX, + "hidden_tap": "result_norm", + "initial_noise": "torch.randn(dtype=vl_embs.dtype)", + "timestep_ids": EXPECTED_TIMESTEP_IDS, + "euler_update": "actions = actions + dt * pred_velocity", + } + + +def _validate_catalog_identity(catalog: Mapping[str, Any]) -> tuple[dict[str, Any], dict[str, Any]]: + variant = get_variant(catalog, VARIANT) + qwen_key = variant.get("qwen_asset") + qwen = catalog.get("shared_assets", {}).get(qwen_key) + if not isinstance(qwen, dict): + raise StarVLAError(f"catalog variant {VARIANT} has no Qwen action asset") + expected_variant = { + "repo_id": OFFICIAL_CHECKPOINT_REPO_ID, + "revision": OFFICIAL_CHECKPOINT_REVISION, + } + for key, expected in expected_variant.items(): + if variant.get(key) != expected: + raise StarVLAError( + f"catalog {VARIANT}.{key} must be {expected!r}, got {variant.get(key)!r}" + ) + expected_checkpoint = { + "path": f"checkpoints/{OFFICIAL_CHECKPOINT_FILENAME}", + "size": OFFICIAL_CHECKPOINT_SIZE, + "sha256": OFFICIAL_CHECKPOINT_SHA256, + } + if variant.get("checkpoint") != expected_checkpoint: + raise StarVLAError("catalog Qwen2.5 GR00T checkpoint identity drifted") + if ( + qwen.get("repo_id") != OFFICIAL_QWEN_REPO_ID + or qwen.get("revision") != OFFICIAL_QWEN_REVISION + ): + raise StarVLAError("catalog Qwen2.5 action-tokenizer identity drifted") + if catalog.get("source_revisions", {}).get("starvla") != OFFICIAL_STARVLA_REVISION: + raise StarVLAError("catalog StarVLA source revision drifted") + return variant, qwen + + +def validate_action_tokenizer_assets(qwen_dir: Path) -> dict[str, Any]: + config = _load_json_object(qwen_dir / "config.json", label="Qwen2.5 action config") + text_config = config.get("text_config") + actual = { + "model_type": config.get("model_type"), + "hidden_size": config.get("hidden_size"), + "text_hidden_size": text_config.get("hidden_size") if isinstance(text_config, dict) else None, + "layer_count": text_config.get("num_hidden_layers") if isinstance(text_config, dict) else None, + "vocab_size": text_config.get("vocab_size") if isinstance(text_config, dict) else None, + } + expected = { + "model_type": BACKBONE, + "hidden_size": EXPECTED_QWEN_HIDDEN_DIM, + "text_hidden_size": EXPECTED_QWEN_HIDDEN_DIM, + "layer_count": EXPECTED_QWEN_LAYER_COUNT, + "vocab_size": ACTION_TOKEN_ID_MAX + 1, + } + if actual != expected: + raise StarVLAError(f"unexpected Qwen2.5 action model config: {actual}") + + token_map = _load_json_object( + qwen_dir / "added_token_id_map.json", + label="Qwen2.5 action token map", + ) + expected_map = { + f"": ACTION_TOKEN_ID_MIN + index + for index in range(ACTION_TOKEN_COUNT) + } + if token_map != expected_map: + raise StarVLAError( + "Qwen2.5 action tokenizer must contain the contiguous " + "2048-token range 151665..153712" + ) + assets, assets_sha256 = _qwen_asset_records(qwen_dir) + return { + "repo_id": OFFICIAL_QWEN_REPO_ID, + "revision": OFFICIAL_QWEN_REVISION, + "action_token_count": ACTION_TOKEN_COUNT, + "action_token_id_min": ACTION_TOKEN_ID_MIN, + "action_token_id_max": ACTION_TOKEN_ID_MAX, + "assets": assets, + "assets_sha256": assets_sha256, + } + + +def _validate_effective_config(config: Mapping[str, Any]) -> None: + try: + framework = config["framework"] + action = framework["action_model"] + diffusion = action["diffusion_model_cfg"] + vla = config["datasets"]["vla_data"] + except (KeyError, TypeError) as exc: + raise StarVLAError("effective Qwen2.5 GR00T config is incomplete") from exc + actual = { + "framework_py": framework.get("framework_py"), + "action_model_type": action.get("action_model_type"), + "action_horizon": action.get("action_horizon"), + "action_dim": action.get("action_dim"), + "state_dim": action.get("state_dim"), + "steps": action.get("num_inference_timesteps"), + "buckets": action.get("num_timestep_buckets"), + "future_tokens": action.get("num_target_vision_tokens"), + "layers": diffusion.get("num_layers"), + "cross_dim": diffusion.get("cross_attention_dim"), + "output_dim": diffusion.get("output_dim"), + "interleave": diffusion.get("interleave_self_attention"), + "obs_image_size": vla.get("obs_image_size"), + "obs": vla.get("obs"), + "data_mix": vla.get("data_mix"), + "cot": vla.get("CoT_prompt"), + } + expected = { + "framework_py": "QwenFM", + "action_model_type": "DiT-B", + "action_horizon": EXPECTED_ACTION_HORIZON, + "action_dim": EXPECTED_ACTION_DIM, + "state_dim": EXPECTED_ACTION_DIM, + "steps": 4, + "buckets": 1000, + "future_tokens": EXPECTED_FUTURE_TOKEN_COUNT, + "layers": EXPECTED_DIT_BLOCK_COUNT, + "cross_dim": EXPECTED_QWEN_HIDDEN_DIM, + "output_dim": EXPECTED_DIT_OUTPUT_DIM, + "interleave": True, + "obs_image_size": None, + "obs": ["image_0"], + "data_mix": "bridge_rt_1", + "cot": EXPECTED_COT_TEMPLATE, + } + if actual != expected: + raise StarVLAError(f"unexpected effective Qwen2.5 GR00T config: {actual}") + + +def groot_normalization_contract( + norm_stats: Mapping[str, Any], unnorm_key: str +) -> dict[str, Any]: + if unnorm_key not in UNNORM_KEYS: + raise StarVLAError( + f"Qwen2.5 GR00T unnorm_key must be one of {list(UNNORM_KEYS)}, " + f"got {unnorm_key!r}" + ) + profile = norm_stats.get(unnorm_key) + action = profile.get("action") if isinstance(profile, Mapping) else None + if not isinstance(action, Mapping): + raise StarVLAError(f"dataset statistics has no {unnorm_key}.action object") + try: + q01 = np.asarray(action["q01"], dtype=np.float32) + q99 = np.asarray(action["q99"], dtype=np.float32) + mask = np.asarray(action["mask"], dtype=np.bool_) + except (KeyError, TypeError, ValueError) as exc: + raise StarVLAError(f"invalid GR00T action statistics for {unnorm_key}: {exc}") from exc + if q01.shape != (7,) or q99.shape != (7,) or mask.shape != (7,): + raise StarVLAError("Qwen2.5 GR00T action statistics must be 7D") + if not np.isfinite(q01).all() or not np.isfinite(q99).all(): + raise StarVLAError("Qwen2.5 GR00T action statistics must be finite") + if np.any(q99[mask] <= q01[mask]): + raise StarVLAError("Qwen2.5 GR00T masked q99 values must exceed q01") + return { + "stats_key": unnorm_key, + "runtime_robot_profile": unnorm_key, + "implementation": "official_PolicyNormProcessor_ComposedModalityTransform", + "q01": q01.tolist(), + "q99": q99.tolist(), + "mask": mask.tolist(), + } + + +def validate_local_inputs( + *, + checkpoint_root: Path, + checkpoint: Path | None, + qwen_model: Path | None, + source_dir: Path, + catalog_path: Path = DEFAULT_CATALOG, +) -> dict[str, Any]: + catalog = load_catalog(catalog_path) + variant, qwen = _validate_catalog_identity(catalog) + checkpoint_root = checkpoint_root.resolve() + policy_dir = checkpoint_root / "sources" / variant["directory"] / variant["revision"] + qwen_dir = checkpoint_root / "sources" / qwen["directory"] / qwen["revision"] + checkpoint_path = policy_dir / variant["checkpoint"]["path"] + if checkpoint is not None and checkpoint.resolve() != checkpoint_path.resolve(): + raise StarVLAError( + f"Qwen2.5 GR00T checkpoint must be the catalog path {checkpoint_path}" + ) + if qwen_model is not None and qwen_model.resolve() != qwen_dir.resolve(): + raise StarVLAError( + f"Qwen2.5 action processor must be the catalog path {qwen_dir}" + ) + + verify_catalog_files(policy_dir, variant) + verify_catalog_files(qwen_dir, qwen) + tokenizer = validate_action_tokenizer_assets(qwen_dir) + source_dir = source_dir.resolve() + revision = _verify_clean_source(source_dir, OFFICIAL_STARVLA_REVISION) + source_probe = verify_source_semantics(source_dir) + + sidecar = Path(f"{checkpoint_path}.aria2") + checkpoint_ready = ( + checkpoint_path.is_file() + and not checkpoint_path.is_symlink() + and not sidecar.exists() + ) + if checkpoint_ready: + verify_checkpoint_file(checkpoint_path, variant) + + config_yaml = policy_dir / "config.yaml" + dataset_statistics = policy_dir / "dataset_statistics.json" + norm_stats = _load_json_object( + dataset_statistics, label="Qwen2.5 GR00T dataset statistics" + ) + if set(norm_stats) != set(UNNORM_KEYS): + raise StarVLAError( + f"unexpected Qwen2.5 GR00T normalization profiles: {sorted(norm_stats)}" + ) + for key in UNNORM_KEYS: + groot_normalization_contract(norm_stats, key) + + try: + import yaml + + config = yaml.safe_load(config_yaml.read_text(encoding="utf-8")) + except (ImportError, OSError, UnicodeError, ValueError) as exc: + raise StarVLAError(f"failed to load Qwen2.5 GR00T config.yaml: {exc}") from exc + if not isinstance(config, dict): + raise StarVLAError("Qwen2.5 GR00T config.yaml root must be an object") + _validate_effective_config(config) + return { + "catalog": catalog, + "catalog_path": catalog_path.resolve(), + "variant": variant, + "qwen": qwen, + "policy_dir": policy_dir.resolve(), + "qwen_dir": qwen_dir.resolve(), + "checkpoint": checkpoint_path.resolve(), + "checkpoint_ready": checkpoint_ready, + "config_yaml": config_yaml.resolve(), + "dataset_statistics": dataset_statistics.resolve(), + "norm_stats": norm_stats, + "config": config, + "source_dir": source_dir, + "source_revision": revision, + "source_probe": source_probe, + "tokenizer": tokenizer, + } + + +def validate_processor_contract(qwen_dir: Path) -> dict[str, Any]: + import transformers + + processor = transformers.AutoProcessor.from_pretrained( + qwen_dir, + local_files_only=True, + trust_remote_code=False, + ) + tokenizer = processor.tokenizer + first = tokenizer("", add_special_tokens=False)["input_ids"] + last = tokenizer("", add_special_tokens=False)["input_ids"] + actual = { + "processor_class": type(processor).__name__, + "image_processor_class": type(processor.image_processor).__name__, + "tokenizer_length": len(tokenizer), + "first_action_token_ids": first, + "last_action_token_ids": last, + "padding_side": tokenizer.padding_side, + } + expected = { + "processor_class": "Qwen2_5_VLProcessor", + "image_processor_class": "Qwen2VLImageProcessorFast", + "tokenizer_length": ACTION_TOKEN_ID_MAX + 1, + "first_action_token_ids": [ACTION_TOKEN_ID_MIN], + "last_action_token_ids": [ACTION_TOKEN_ID_MAX], + "padding_side": "right", + } + # The official wrapper switches padding to left immediately after loading. + if actual != expected: + raise StarVLAError(f"unexpected Qwen2.5 action processor contract: {actual}") + return {**actual, "wrapper_padding_side": "left"} + + +def load_official_framework(paths: Mapping[str, Any], *, device: str) -> tuple[Any, dict[str, Any]]: + import torch + import transformers + + if not paths["checkpoint_ready"]: + raise StarVLAError( + f"official Qwen2.5 GR00T checkpoint is absent or incomplete: {paths['checkpoint']}" + ) + source_dir = Path(paths["source_dir"]) + if any(name == "starVLA" or name.startswith("starVLA.") for name in sys.modules): + raise StarVLAError("starVLA was imported before pinned-source verification") + sys.path.insert(0, str(source_dir)) + try: + from starVLA.model.framework import base_framework, share_tools + from starVLA.model.framework.VLM4A import QwenGR00T + + for module in (base_framework, share_tools, QwenGR00T): + _assert_module_origin(module, source_dir) + config, norm_stats = share_tools.read_mode_config(str(paths["checkpoint"])) + _validate_effective_config(config) + with _official_qwen25_alias(Path(paths["qwen_dir"])) as qwen_alias: + config = base_framework.merge_config_overrides( + config, + [ + f"framework.qwenvl.base_vlm={qwen_alias}", + "framework.qwenvl.attn_implementation=sdpa", + ], + ) + cfg = share_tools.dict_to_namespace(config) + cfg.trainer.pretrained_checkpoint = None + with _config_only_qwen25_bootstrap( + torch, transformers, Path(paths["qwen_dir"]) + ): + framework = QwenGR00T.Qwen_GR00T(cfg) + + try: + state = torch.load( + paths["checkpoint"], map_location="cpu", mmap=True, weights_only=True + ) + except TypeError: + state = torch.load(paths["checkpoint"], map_location="cpu", weights_only=True) + if not isinstance(state, Mapping) or not state: + raise StarVLAError("official Qwen2.5 GR00T checkpoint has no state_dict") + framework.load_state_dict(state, strict=True) + del state + gc.collect() + framework.norm_stats = norm_stats + + action_model = framework.action_model + if type(framework).__name__ != "Qwen_GR00T": + raise StarVLAError(f"unexpected official framework class: {type(framework).__name__}") + if int(framework.action_horizon) != EXPECTED_ACTION_HORIZON: + raise StarVLAError("official Qwen2.5 GR00T action horizon changed") + if len(action_model.model.transformer_blocks) != EXPECTED_DIT_BLOCK_COUNT: + raise StarVLAError("official Qwen2.5 GR00T DiT block count changed") + hidden_size = int(framework.qwen_vl_interface.model.config.hidden_size) + if hidden_size != EXPECTED_QWEN_HIDDEN_DIM: + raise StarVLAError(f"unexpected Qwen2.5 hidden size: {hidden_size}") + tokenizer = framework.qwen_vl_interface.processor.tokenizer + if ( + len(tokenizer) != ACTION_TOKEN_ID_MAX + 1 + or tokenizer.convert_tokens_to_ids("") != ACTION_TOKEN_ID_MIN + or tokenizer.convert_tokens_to_ids("") != ACTION_TOKEN_ID_MAX + ): + raise StarVLAError("official framework did not load the pinned action tokenizer") + + qwen_dtypes = {parameter.dtype for parameter in framework.qwen_vl_interface.parameters()} + policy_dtypes = {parameter.dtype for parameter in action_model.parameters()} + if qwen_dtypes != {torch.bfloat16} or policy_dtypes != {torch.float32}: + raise StarVLAError( + "official Qwen2.5 GR00T dtype boundary changed: " + f"qwen={qwen_dtypes}, policy={policy_dtypes}" + ) + return framework.to(device).eval(), config + finally: + if sys.path and sys.path[0] == str(source_dir): + del sys.path[0] + + +def _first_tensor(value: Any) -> Any: + if isinstance(value, (tuple, list)): + if not value: + raise StarVLAError("decoder layer returned an empty tuple") + return value[0] + return value + + +def run_official_forward( + framework: Any, + *, + images: Sequence[Any], + task: str, + seed: int = SEED, +) -> dict[str, Any]: + """Run the official framework and capture its result_norm/fixed-noise boundary.""" + + import torch + + captures: dict[str, Any] = {} + qwen = framework.qwen_vl_interface + action_model = framework.action_model + language_model = qwen.model.model.language_model + handles = [] + original_build = qwen.build_qwenvl_inputs + original_policy = action_model.predict_action + + def capture_build(*args: Any, **kwargs: Any): + batch_images = kwargs.get("images", args[0] if args else None) + instructions = kwargs.get("instructions", args[1] if len(args) > 1 else None) + captures["processed_images"] = list(batch_images[0]) + captures["framework_instructions"] = list(instructions) + output = original_build(*args, **kwargs) + captures["qwen_inputs"] = { + key: value.detach() + for key, value in output.items() + if isinstance(value, torch.Tensor) + } + return output + + def capture_outer(_module: Any, _inputs: Any, output: Any): + hidden = getattr(output, "hidden_states", None) + if hidden is None or len(hidden) != EXPECTED_QWEN_LAYER_COUNT + 1: + raise StarVLAError( + "official Qwen2.5 outer output did not expose 37 hidden tuple entries" + ) + captures["outer_final"] = hidden[EXPECTED_HIDDEN_TUPLE_INDEX].detach().clone() + + def capture_policy(*args: Any, **kwargs: Any): + conditioning = args[0] if args else kwargs.get("vl_embs") + state = args[1] if len(args) > 1 else kwargs.get("state") + mask = kwargs.get("encoder_attention_mask", args[2] if len(args) > 2 else None) + if state is not None: + raise StarVLAError("official Qwen2.5 GR00T unexpectedly used state") + captures["policy_conditioning"] = conditioning.detach().clone() + captures["policy_attention_mask"] = mask.detach().clone() + original_randn = torch.randn + + def capture_randn(*randn_args: Any, **randn_kwargs: Any): + value = original_randn(*randn_args, **randn_kwargs) + if "initial_noise" in captures: + raise StarVLAError("official Qwen2.5 GR00T sampled noise more than once") + captures["initial_noise"] = value.detach().clone() + return value + + torch.randn = capture_randn + try: + output = original_policy(*args, **kwargs) + finally: + torch.randn = original_randn + captures["raw_policy"] = output.detach().clone() + return output + + handles.append( + language_model.layers[-1].register_forward_hook( + lambda _m, _i, output: captures.__setitem__( + "raw_l_out_35", _first_tensor(output).detach().clone() + ) + ) + ) + handles.append( + language_model.norm.register_forward_hook( + lambda _m, _i, output: captures.__setitem__( + "result_norm", output.detach().clone() + ) + ) + ) + handles.append(qwen.model.register_forward_hook(capture_outer)) + qwen.build_qwenvl_inputs = capture_build + action_model.predict_action = capture_policy + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + try: + result = framework.predict_action( + examples=[{"image": list(images), "lang": task}] + ) + finally: + for handle in handles: + handle.remove() + qwen.build_qwenvl_inputs = original_build + action_model.predict_action = original_policy + + required = { + "processed_images", + "framework_instructions", + "qwen_inputs", + "raw_l_out_35", + "result_norm", + "outer_final", + "policy_conditioning", + "policy_attention_mask", + "initial_noise", + "raw_policy", + } + missing = sorted(required - set(captures)) + if missing: + raise StarVLAError(f"Qwen2.5 GR00T instrumentation missed: {missing}") + if captures["framework_instructions"] != [task]: + raise StarVLAError("official framework changed the input instruction") + if len(captures["processed_images"]) != len(images) or any( + actual.mode != expected.mode + or actual.size != expected.size + or actual.tobytes() != expected.tobytes() + for actual, expected in zip(captures["processed_images"], images, strict=True) + ): + raise StarVLAError("official Qwen2.5 GR00T unexpectedly pre-resized the image") + if not torch.equal(captures["outer_final"], captures["result_norm"]): + raise StarVLAError("Qwen2.5 hidden tuple entry 36 is not result_norm") + if torch.equal(captures["outer_final"], captures["raw_l_out_35"]): + raise StarVLAError("Qwen2.5 result_norm unexpectedly equals raw l_out-35") + if not torch.equal(captures["policy_conditioning"], captures["result_norm"]): + raise StarVLAError("Qwen2.5 GR00T policy did not receive result_norm") + mask = captures["qwen_inputs"].get("attention_mask") + if mask is None or not torch.equal( + captures["policy_attention_mask"], mask.to(dtype=torch.bool) + ): + raise StarVLAError("Qwen2.5 GR00T policy mask is not the full boolean Qwen mask") + if captures["policy_attention_mask"].dtype != torch.bool: + raise StarVLAError("Qwen2.5 GR00T policy attention mask must be bool") + if captures["result_norm"].dtype != torch.bfloat16: + raise StarVLAError("Qwen2.5 result_norm boundary must be BF16") + if captures["initial_noise"].dtype != torch.bfloat16: + raise StarVLAError("Qwen2.5 GR00T initial noise must be sampled as BF16") + if tuple(captures["initial_noise"].shape) != (1, 16, 7): + raise StarVLAError("Qwen2.5 GR00T initial noise shape must be [1,16,7]") + + normalized = np.asarray(result.get("normalized_actions"), dtype=np.float32) + raw_policy, _ = _tensor_to_array(captures["raw_policy"]) + if normalized.shape != (1, EXPECTED_ACTION_HORIZON, EXPECTED_ACTION_DIM): + raise StarVLAError(f"unexpected Qwen2.5 GR00T output shape: {normalized.shape}") + if normalized.shape != raw_policy.shape or not np.array_equal(normalized, raw_policy): + raise StarVLAError("normalized actions differ from captured GR00T policy output") + if not np.isfinite(normalized).all(): + raise StarVLAError("official Qwen2.5 GR00T produced non-finite actions") + captures["normalized_actions"] = np.ascontiguousarray(normalized) + return captures + + +def _load_images(image_paths: Iterable[Path]) -> tuple[list[Any], list[dict[str, Any]]]: + from PIL import Image + + images: list[Any] = [] + records: list[dict[str, Any]] = [] + for path in image_paths: + path = path.resolve() + _ensure_regular_file(path, label="Qwen2.5 GR00T input image") + try: + with Image.open(path) as opened: + opened.load() + image = opened.convert("RGB") + except (OSError, ValueError) as exc: + raise StarVLAError(f"failed to decode input image {path}: {exc}") from exc + images.append(image) + records.append( + { + "source_path": str(path), + "source_size": path.stat().st_size, + "source_sha256": sha256_file(path), + "decoded_mode": image.mode, + "decoded_size": list(image.size), + "decoded_pixel_sha256": _image_pixel_sha256(image), + } + ) + if len(images) != 1: + raise StarVLAError("official Qwen2.5 GR00T oracle requires exactly one image") + return images, records + + +def _render_model_prompt(framework: Any, images: Sequence[Any], task: str) -> str: + model_instruction = EXPECTED_COT_TEMPLATE.replace("{instruction}", task) + messages = [{ + "role": "user", + "content": [ + *({"type": "image", "image": image} for image in images), + {"type": "text", "text": model_instruction}, + ], + }] + rendered = framework.qwen_vl_interface.processor.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True + ) + if not isinstance(rendered, str): + raise StarVLAError("Qwen2.5 processor returned a non-string prompt") + return rendered + + +def _build_arrays( + captures: Mapping[str, Any], unnormalized: np.ndarray +) -> tuple[dict[str, np.ndarray], dict[str, Any]]: + required_inputs = ("input_ids", "attention_mask", "image_grid_thw") + missing = [name for name in required_inputs if name not in captures["qwen_inputs"]] + if missing: + raise StarVLAError(f"Qwen2.5 processor outputs are missing: {missing}") + values = { + "input_ids": captures["qwen_inputs"]["input_ids"], + "attention_mask": captures["qwen_inputs"]["attention_mask"], + "image_grid_thw": captures["qwen_inputs"]["image_grid_thw"], + "raw_l_out_35_diagnostic": captures["raw_l_out_35"], + "result_norm": captures["result_norm"], + "initial_noise": captures["initial_noise"], + "normalized_actions": captures["normalized_actions"], + "unnormalized_actions": np.ascontiguousarray(unnormalized, dtype=np.float32), + } + arrays: dict[str, np.ndarray] = {} + records: dict[str, Any] = {} + for name, value in values.items(): + if isinstance(value, np.ndarray): + array = np.ascontiguousarray(value) + source_dtype = None + else: + array, source_dtype = _tensor_to_array(value) + arrays[name] = array + records[name] = _array_record(array, source_dtype=source_dtype) + + token_count = arrays["input_ids"].shape[1] + expected_shapes = { + "input_ids": (1, token_count), + "attention_mask": (1, token_count), + "image_grid_thw": (1, 3), + "raw_l_out_35_diagnostic": (1, token_count, EXPECTED_QWEN_HIDDEN_DIM), + "result_norm": (1, token_count, EXPECTED_QWEN_HIDDEN_DIM), + "initial_noise": (1, EXPECTED_ACTION_HORIZON, EXPECTED_ACTION_DIM), + "normalized_actions": (1, EXPECTED_ACTION_HORIZON, EXPECTED_ACTION_DIM), + "unnormalized_actions": (1, EXPECTED_ACTION_HORIZON, EXPECTED_ACTION_DIM), + } + for name, expected in expected_shapes.items(): + if arrays[name].shape != expected: + raise StarVLAError( + f"Qwen2.5 GR00T {name} shape must be {expected}, got {arrays[name].shape}" + ) + return arrays, records + + +def write_golden( + *, + output_dir: Path, + paths: Mapping[str, Any], + framework: Any, + image_paths: Sequence[Path], + source_image_records: Sequence[Mapping[str, Any]], + task: str, + unnorm_key: str, + captures: Mapping[str, Any], + unnormalized: np.ndarray, +) -> Path: + import torch + import transformers + + output_dir = output_dir.resolve() + if output_dir.exists(): + raise StarVLAError(f"golden output directory already exists: {output_dir}") + output_dir.parent.mkdir(parents=True, exist_ok=True) + arrays, records = _build_arrays(captures, unnormalized) + model_instruction = EXPECTED_COT_TEMPLATE.replace("{instruction}", task) + rendered_prompt = _render_model_prompt( + framework, captures["processed_images"], task + ) + identity = { + "schema_version": SCHEMA_VERSION, + "kind": GOLDEN_KIND, + "checkpoint_sha256": OFFICIAL_CHECKPOINT_SHA256, + "starvla_revision": paths["source_revision"], + "qwen_revision": OFFICIAL_QWEN_REVISION, + "task": task, + "unnorm_key": unnorm_key, + "images": [record["source_sha256"] for record in source_image_records], + "initial_noise_sha256": records["initial_noise"]["sha256"], + } + golden_id = _sha256_bytes(_canonical_json(identity)) + + with tempfile.TemporaryDirectory( + prefix=f".{output_dir.name}.", dir=output_dir.parent + ) as temporary: + staging = Path(temporary) + inputs_dir = staging / "inputs" + inputs_dir.mkdir() + image_records: list[dict[str, Any]] = [] + for index, (source_path, source_record) in enumerate( + zip(image_paths, source_image_records, strict=True) + ): + suffix = source_path.suffix.lower() or ".img" + artifact = inputs_dir / f"image-{index:02d}{suffix}" + shutil.copyfile(source_path, artifact) + image_records.append({ + **source_record, + "artifact": artifact.relative_to(staging).as_posix(), + "artifact_size": artifact.stat().st_size, + "artifact_sha256": sha256_file(artifact), + }) + + tensors_path = staging / "tensors.npz" + np.savez(tensors_path, **arrays) + noise_bytes = np.ascontiguousarray( + arrays["initial_noise"], dtype=" dict[str, Any]: + return { + "schema_version": SCHEMA_VERSION, + "kind": "starvla_qwen25_groot_preflight", + "variant": VARIANT, + "model_type": MODEL_TYPE, + "backbone": BACKBONE, + "checkpoint": str(paths["checkpoint"]), + "checkpoint_ready": paths["checkpoint_ready"], + "expected_checkpoint": { + "bundle_uuid": official_bundle_uuid( + paths["variant"], paths["catalog"] + ), + "repo_id": OFFICIAL_CHECKPOINT_REPO_ID, + "revision": OFFICIAL_CHECKPOINT_REVISION, + "filename": OFFICIAL_CHECKPOINT_FILENAME, + "size": OFFICIAL_CHECKPOINT_SIZE, + "sha256": OFFICIAL_CHECKPOINT_SHA256, + }, + "qwen": { + **paths["tokenizer"], + "processor": dict(processor), + }, + "conditioning": { + "hidden_tuple_index": EXPECTED_HIDDEN_TUPLE_INDEX, + "hidden_tap_name": "result_norm", + "hidden_size": EXPECTED_QWEN_HIDDEN_DIM, + }, + "action": { + "shape": [1, EXPECTED_ACTION_HORIZON, EXPECTED_ACTION_DIM], + "initial_noise_dtype": "bfloat16", + "timestep_ids": EXPECTED_TIMESTEP_IDS, + }, + "action_gate": { + "reference": "local_official_python_pt", + "metric": "full_tensor_global_relative_l2", + "operator": "<=", + "limit": ACTION_RELATIVE_L2_LIMIT, + "required_outputs": ["normalized_actions", "unnormalized_actions"], + }, + "source_probe": paths["source_probe"], + "effective_config_valid": True, + "golden_created": False, + } + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--checkpoint-root", type=Path, default=Path("ckpts/starvla")) + parser.add_argument("--checkpoint", type=Path) + parser.add_argument("--qwen-model", type=Path) + parser.add_argument( + "--starvla-source", + type=Path, + default=Path("ckpts/starvla/source/starvla"), + ) + parser.add_argument("--catalog", type=Path, default=DEFAULT_CATALOG) + parser.add_argument("--image", action="append", default=[], type=Path) + parser.add_argument("--task", default="grab the block.") + parser.add_argument("--unnorm-key", choices=UNNORM_KEYS, default="oxe_bridge") + parser.add_argument("--device", default="cuda:0") + parser.add_argument( + "--output-dir", + type=Path, + default=Path("goldens/starvla/qwen25-groot/bridge-grab-block"), + ) + parser.add_argument("--preflight", "--preflight-only", action="store_true") + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + args = build_parser().parse_args(argv) + try: + _require_isolated_python() + import torch + import transformers + + validate_runtime_versions( + torch_version=torch.__version__, + torchvision_version=_distribution_version("torchvision"), + transformers_version=transformers.__version__, + numpy_version=np.__version__, + ) + if _distribution_version("qwen-vl-utils") != EXPECTED_QWEN_VL_UTILS_VERSION: + raise StarVLAError( + "qwen-vl-utils must be " + f"{EXPECTED_QWEN_VL_UTILS_VERSION} for the official oracle" + ) + _configure_determinism(torch, seed=SEED, device=args.device) + paths = validate_local_inputs( + checkpoint_root=args.checkpoint_root, + checkpoint=args.checkpoint, + qwen_model=args.qwen_model, + source_dir=args.starvla_source, + catalog_path=args.catalog, + ) + processor = validate_processor_contract(Path(paths["qwen_dir"])) + if args.preflight: + print(json.dumps( + _preflight_record(paths, processor), + allow_nan=False, + indent=2, + sort_keys=True, + )) + return 0 + if not paths["checkpoint_ready"]: + raise StarVLAError( + f"official Qwen2.5 GR00T checkpoint is not ready: {paths['checkpoint']}" + ) + if len(args.image) != 1: + raise StarVLAError("exactly one --image is required") + images, image_records = _load_images(args.image) + framework, _config = load_official_framework(paths, device=args.device) + captures = run_official_forward( + framework, images=images, task=args.task, seed=SEED + ) + + source_dir = Path(paths["source_dir"]) + sys.path.insert(0, str(source_dir)) + try: + from deployment.model_server import policy_norm_processor + + _assert_module_origin(policy_norm_processor, source_dir) + normalizer = policy_norm_processor.PolicyNormProcessor( + str(paths["checkpoint"]), unnorm_key=args.unnorm_key + ) + unnormalized = np.stack([ + normalizer.unapply_actions(captures["normalized_actions"][0]) + ]).astype(np.float32, copy=False) + finally: + if sys.path and sys.path[0] == str(source_dir): + del sys.path[0] + if unnormalized.shape != (1, EXPECTED_ACTION_HORIZON, EXPECTED_ACTION_DIM): + raise StarVLAError( + f"official Qwen2.5 GR00T unnormalized shape changed: {unnormalized.shape}" + ) + if not np.isfinite(unnormalized).all(): + raise StarVLAError("official Qwen2.5 GR00T unnormalized actions are non-finite") + + manifest = write_golden( + output_dir=args.output_dir, + paths=paths, + framework=framework, + image_paths=args.image, + source_image_records=image_records, + task=args.task, + unnorm_key=args.unnorm_key, + captures=captures, + unnormalized=unnormalized, + ) + print(manifest) + return 0 + except (StarVLAError, OSError, ValueError, KeyError, RuntimeError) as exc: + print(f"error: {exc}", file=sys.stderr) + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) From 11cc5bf05f31319fbd84cf13d4fb75149eb463e8 Mon Sep 17 00:00:00 2001 From: JJJYmmm <1650675829@qq.com> Date: Mon, 10 Aug 2026 12:42:56 +0800 Subject: [PATCH 08/11] starvla: add Qwen2.5-VL FAST policy --- src/models/starvla/fast_codec.cpp | 852 +++++++++ src/models/starvla/fast_codec.h | 90 + src/models/starvla/fast_policy.cpp | 459 +++++ src/models/starvla/fast_policy.h | 79 + tests/starvla/fast_codec_test.cpp | 194 +++ tests/starvla/fast_runtime_test.cpp | 137 ++ tests/starvla/test_starvla_qwen25_fast.py | 111 ++ .../starvla/convert_starvla_qwen25_fast.py | 1542 +++++++++++++++++ .../generate_starvla_qwen25_fast_golden.py | 765 ++++++++ 9 files changed, 4229 insertions(+) create mode 100644 src/models/starvla/fast_codec.cpp create mode 100644 src/models/starvla/fast_codec.h create mode 100644 src/models/starvla/fast_policy.cpp create mode 100644 src/models/starvla/fast_policy.h create mode 100644 tests/starvla/fast_codec_test.cpp create mode 100644 tests/starvla/fast_runtime_test.cpp create mode 100644 tests/starvla/test_starvla_qwen25_fast.py create mode 100644 tools/hf2gguf/starvla/convert_starvla_qwen25_fast.py create mode 100644 tools/hf2gguf/starvla/generate_starvla_qwen25_fast_golden.py diff --git a/src/models/starvla/fast_codec.cpp b/src/models/starvla/fast_codec.cpp new file mode 100644 index 0000000..3a01ac9 --- /dev/null +++ b/src/models/starvla/fast_codec.cpp @@ -0,0 +1,852 @@ +#include "models/starvla/fast_codec.h" + +#include "nlohmann/json.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace robotcpp::starvla { +namespace { + +using Json = nlohmann::json; + +constexpr size_t kMaximumJsonBytes = 16U * 1024U * 1024U; +constexpr size_t kOfficialVocabSize = 2048U; +constexpr size_t kMaximumVocabSize = 65536U; +constexpr size_t kMaximumTimeHorizon = 1024U; +constexpr size_t kMaximumActionDim = 1024U; +constexpr size_t kMaximumBatchSize = 1024U; +constexpr size_t kMaximumTokenSequence = 4096U; +constexpr size_t kMaximumGeneratedSequence = 2048U; +constexpr size_t kMaximumDecodedBytes = 1024U * 1024U; +constexpr size_t kMaximumOutputScalars = 16U * 1024U * 1024U; +constexpr uint64_t kMaximumIdctMultiplyAdds = 64ULL * 1024ULL * 1024ULL; +constexpr const char * kActionTokenPrefix = " kMaximumJsonBytes) { + error = "StarVLA FAST JSON asset exceeds the 16 MiB limit: " + path.string(); + return false; + } + + std::ifstream stream(path, std::ios::binary); + if (!stream) { + error = "cannot open StarVLA FAST JSON asset: " + path.string(); + return false; + } + std::string contents; + contents.reserve(static_cast(size)); + std::array buffer{}; + while (stream) { + stream.read(buffer.data(), static_cast(buffer.size())); + const std::streamsize count = stream.gcount(); + if (count <= 0) { + continue; + } + const size_t chunk_size = static_cast(count); + if (contents.size() > kMaximumJsonBytes - chunk_size) { + error = "StarVLA FAST JSON asset exceeds the 16 MiB limit while reading: " + + path.string(); + return false; + } + contents.append(buffer.data(), chunk_size); + } + if (!stream.eof() || stream.bad()) { + error = "cannot read StarVLA FAST JSON asset: " + path.string(); + return false; + } + output = Json::parse(contents, nullptr, false); + if (output.is_discarded()) { + error = "cannot parse StarVLA FAST JSON asset: " + path.string(); + return false; + } + return true; +} + +bool decode_utf8_strict(const std::string & input, std::vector & output, + std::string & error) { + output.clear(); + for (size_t i = 0; i < input.size();) { + const uint8_t first = static_cast(input[i]); + uint32_t value = 0; + size_t length = 0; + if (first <= 0x7fU) { + value = first; + length = 1; + } else if (first >= 0xc2U && first <= 0xdfU) { + value = first & 0x1fU; + length = 2; + } else if (first >= 0xe0U && first <= 0xefU) { + value = first & 0x0fU; + length = 3; + } else if (first >= 0xf0U && first <= 0xf4U) { + value = first & 0x07U; + length = 4; + } else { + error = "StarVLA FAST tokenizer vocabulary contains invalid UTF-8"; + return false; + } + if (i + length > input.size()) { + error = "StarVLA FAST tokenizer vocabulary contains truncated UTF-8"; + return false; + } + for (size_t j = 1; j < length; ++j) { + const uint8_t continuation = static_cast(input[i + j]); + if ((continuation & 0xc0U) != 0x80U) { + error = "StarVLA FAST tokenizer vocabulary contains invalid UTF-8 continuation"; + return false; + } + value = (value << 6U) | (continuation & 0x3fU); + } + const bool overlong = (length == 2 && value < 0x80U) || + (length == 3 && value < 0x800U) || + (length == 4 && value < 0x10000U); + if (overlong || value > 0x10ffffU || (value >= 0xd800U && value <= 0xdfffU)) { + error = "StarVLA FAST tokenizer vocabulary contains a non-scalar UTF-8 value"; + return false; + } + output.push_back(value); + i += length; + } + return true; +} + +std::unordered_map byte_level_inverse_alphabet() { + std::unordered_map result; + std::unordered_set direct; + for (int value = 0x21; value <= 0x7e; ++value) { + direct.insert(value); + result.emplace(static_cast(value), static_cast(value)); + } + for (int value = 0xa1; value <= 0xac; ++value) { + direct.insert(value); + result.emplace(static_cast(value), static_cast(value)); + } + for (int value = 0xae; value <= 0xff; ++value) { + direct.insert(value); + result.emplace(static_cast(value), static_cast(value)); + } + uint32_t extra = 0; + for (int value = 0; value <= 0xff; ++value) { + if (direct.count(value) == 0) { + result.emplace(256U + extra, static_cast(value)); + ++extra; + } + } + return result; +} + +bool compile_token_bytes(const std::vector & vocab_by_id, + std::vector> & token_bytes, + std::string & error) { + const auto inverse_alphabet = byte_level_inverse_alphabet(); + token_bytes.clear(); + token_bytes.reserve(vocab_by_id.size()); + for (size_t token_id = 0; token_id < vocab_by_id.size(); ++token_id) { + if (vocab_by_id[token_id].empty()) { + error = "StarVLA FAST tokenizer has an empty vocabulary piece at ID " + + std::to_string(token_id); + return false; + } + std::vector piece_codepoints; + if (!decode_utf8_strict(vocab_by_id[token_id], piece_codepoints, error)) { + error += " at token ID " + std::to_string(token_id); + return false; + } + std::vector bytes; + bytes.reserve(piece_codepoints.size()); + for (uint32_t codepoint : piece_codepoints) { + const auto found = inverse_alphabet.find(codepoint); + if (found == inverse_alphabet.end()) { + error = "StarVLA FAST tokenizer piece contains a code point outside the ByteLevel " + "alphabet at ID " + + std::to_string(token_id); + return false; + } + bytes.push_back(found->second); + } + token_bytes.push_back(std::move(bytes)); + } + return true; +} + +void decode_utf8_lossy(const std::vector & input, std::vector & output) { + output.clear(); + for (size_t i = 0; i < input.size();) { + const uint8_t first = input[i]; + if (first <= 0x7fU) { + output.push_back(first); + ++i; + continue; + } + + size_t length = 0; + uint32_t value = 0; + if (first >= 0xc2U && first <= 0xdfU) { + length = 2; + value = first & 0x1fU; + } else if (first >= 0xe0U && first <= 0xefU) { + length = 3; + value = first & 0x0fU; + } else if (first >= 0xf0U && first <= 0xf4U) { + length = 4; + value = first & 0x07U; + } else { + output.push_back(0xfffdU); + ++i; + continue; + } + + if (i + 1 >= input.size()) { + output.push_back(0xfffdU); + break; + } + const uint8_t second = input[i + 1]; + const bool second_is_continuation = (second & 0xc0U) == 0x80U; + const bool second_in_scalar_range = + !(first == 0xe0U && second < 0xa0U) && + !(first == 0xedU && second > 0x9fU) && + !(first == 0xf0U && second < 0x90U) && + !(first == 0xf4U && second > 0x8fU); + if (!second_is_continuation || !second_in_scalar_range) { + output.push_back(0xfffdU); + ++i; + continue; + } + value = (value << 6U) | (second & 0x3fU); + + bool invalid = false; + size_t consumed_prefix = 2; + for (size_t j = 2; j < length; ++j) { + if (i + j >= input.size()) { + output.push_back(0xfffdU); + i = input.size(); + invalid = true; + break; + } + const uint8_t continuation = input[i + j]; + if ((continuation & 0xc0U) != 0x80U) { + output.push_back(0xfffdU); + i += consumed_prefix; + invalid = true; + break; + } + value = (value << 6U) | (continuation & 0x3fU); + ++consumed_prefix; + } + if (invalid) { + continue; + } + output.push_back(value); + i += length; + } +} + +bool parse_positive_size(const Json & value, const char * name, size_t & output, + std::string & error) { + if (!value.is_number_integer()) { + error = std::string("StarVLA FAST ") + name + " must be an integer"; + return false; + } + try { + const int64_t parsed = value.get(); + if (parsed <= 0 || static_cast(parsed) > + static_cast(std::numeric_limits::max())) { + error = std::string("StarVLA FAST ") + name + " is out of range"; + return false; + } + output = static_cast(parsed); + return true; + } catch (const std::exception &) { + error = std::string("StarVLA FAST ") + name + " is out of range"; + return false; + } +} + +bool parse_processor_config(const Json & json, size_t time_horizon_override, + size_t action_dim_override, FastCodecConfig & config, + std::string & error) { + if (!json.is_object() || !json.contains("processor_class") || + json["processor_class"] != "UniversalActionProcessor" || !json.contains("scale") || + !json["scale"].is_number() || !json.contains("vocab_size") || + !json.contains("min_token") || !json["min_token"].is_number_integer()) { + error = "StarVLA FAST processor_config.json has an incompatible schema"; + return false; + } + config.scale = json["scale"].get(); + if (!std::isfinite(config.scale) || config.scale == 0.0) { + error = "StarVLA FAST processor scale must be finite and non-zero"; + return false; + } + if (!parse_positive_size(json["vocab_size"], "vocab_size", config.vocab_size, error)) { + return false; + } + if (config.vocab_size != kOfficialVocabSize) { + error = "StarVLA FAST pinned vocabulary must contain exactly 2048 tokens"; + return false; + } + try { + if (json["min_token"].is_number_unsigned()) { + const uint64_t min_token = json["min_token"].get(); + if (min_token > static_cast(std::numeric_limits::max())) { + error = "StarVLA FAST min_token is out of int32 range"; + return false; + } + config.min_token = static_cast(min_token); + } else { + const int64_t min_token = json["min_token"].get(); + if (min_token < std::numeric_limits::min() || + min_token > std::numeric_limits::max()) { + error = "StarVLA FAST min_token is out of int32 range"; + return false; + } + config.min_token = static_cast(min_token); + } + } catch (const std::exception &) { + error = "StarVLA FAST min_token is out of int32 range"; + return false; + } + + auto choose_dimension = [&](const char * name, size_t override_value, size_t & target) { + if (override_value != 0) { + target = override_value; + return true; + } + if (!json.contains(name) || json[name].is_null()) { + error = std::string("StarVLA FAST ") + name + + " is absent; pass the policy dimension explicitly"; + return false; + } + return parse_positive_size(json[name], name, target, error); + }; + return choose_dimension("time_horizon", time_horizon_override, config.time_horizon) && + choose_dimension("action_dim", action_dim_override, config.action_dim); +} + +bool parse_tokenizer_vocab(const Json & json, size_t expected_vocab_size, + std::vector & vocab_by_id, std::string & error) { + if (!json.is_object() || !json.contains("version") || json["version"] != "1.0" || + !json.contains("added_tokens") || !json["added_tokens"].is_array() || + !json["added_tokens"].empty() || !json.contains("decoder") || + !json["decoder"].is_object() || !json["decoder"].contains("type") || + json["decoder"]["type"] != "ByteLevel" || !json.contains("model") || + !json["model"].is_object() || !json["model"].contains("type") || + json["model"]["type"] != "BPE" || !json["model"].contains("vocab") || + !json["model"]["vocab"].is_object()) { + error = "StarVLA FAST tokenizer.json is not the required ByteLevel BPE schema"; + return false; + } + const Json & decoder = json["decoder"]; + if (decoder.size() != 4 || !decoder.contains("add_prefix_space") || + decoder["add_prefix_space"] != true || !decoder.contains("trim_offsets") || + decoder["trim_offsets"] != true || !decoder.contains("use_regex") || + decoder["use_regex"] != true) { + error = "StarVLA FAST tokenizer.json has an incompatible ByteLevel decoder contract"; + return false; + } + const Json & vocab = json["model"]["vocab"]; + if (vocab.size() != expected_vocab_size) { + error = "StarVLA FAST tokenizer vocabulary size does not match processor_config.json"; + return false; + } + vocab_by_id.assign(expected_vocab_size, std::string()); + std::vector seen(expected_vocab_size, false); + for (auto iterator = vocab.begin(); iterator != vocab.end(); ++iterator) { + if (!iterator.value().is_number_integer()) { + error = "StarVLA FAST tokenizer vocabulary ID is not an integer"; + return false; + } + int64_t token_id = -1; + try { + token_id = iterator.value().get(); + } catch (const std::exception &) { + error = "StarVLA FAST tokenizer vocabulary ID is out of range"; + return false; + } + if (token_id < 0 || static_cast(token_id) >= expected_vocab_size || + seen[static_cast(token_id)]) { + error = "StarVLA FAST tokenizer vocabulary IDs are not a bijection"; + return false; + } + seen[static_cast(token_id)] = true; + vocab_by_id[static_cast(token_id)] = iterator.key(); + } + return true; +} + +bool parse_action_index(const std::string & value, size_t & index) { + const std::string prefix(kActionTokenPrefix); + if (value.size() <= prefix.size() + 1 || value.compare(0, prefix.size(), prefix) != 0 || + value.back() != '>') { + return false; + } + const std::string digits = value.substr(prefix.size(), value.size() - prefix.size() - 1); + if (digits.empty() || (digits.size() > 1 && digits.front() == '0')) { + return false; + } + size_t parsed = 0; + for (char character : digits) { + if (character < '0' || character > '9') { + return false; + } + const size_t digit = static_cast(character - '0'); + if (parsed > (std::numeric_limits::max() - digit) / 10U) { + return false; + } + parsed = parsed * 10U + digit; + } + index = parsed; + return true; +} + +bool parse_action_map(const Json & json, size_t vocab_size, + std::vector & fast_to_vlm, std::string & error) { + if (!json.is_object() || json.size() != vocab_size) { + error = "StarVLA FAST action-token map must contain exactly one entry per FAST token"; + return false; + } + fast_to_vlm.assign(vocab_size, -1); + std::vector seen(vocab_size, false); + std::unordered_set vlm_ids; + for (auto iterator = json.begin(); iterator != json.end(); ++iterator) { + size_t fast_id = 0; + if (!parse_action_index(iterator.key(), fast_id) || fast_id >= vocab_size || seen[fast_id]) { + error = "StarVLA FAST action-token map has a malformed or duplicate token name"; + return false; + } + if (!iterator.value().is_number_integer()) { + error = "StarVLA FAST action-token map contains a non-integer VLM ID"; + return false; + } + int64_t vlm_id = -1; + try { + vlm_id = iterator.value().get(); + } catch (const std::exception &) { + error = "StarVLA FAST action-token VLM ID is out of range"; + return false; + } + if (vlm_id < 0 || vlm_id > std::numeric_limits::max() || + !vlm_ids.insert(static_cast(vlm_id)).second) { + error = "StarVLA FAST action-token VLM IDs must be unique non-negative int32 values"; + return false; + } + seen[fast_id] = true; + fast_to_vlm[fast_id] = static_cast(vlm_id); + } + return true; +} + +bool checked_action_count(const FastCodecConfig & config, size_t batch_size, + size_t & per_sample, size_t & total, std::string & error) { + if (config.vocab_size > kMaximumVocabSize || config.time_horizon > kMaximumTimeHorizon || + config.action_dim > kMaximumActionDim) { + error = "StarVLA FAST codec dimensions exceed the runtime safety limits"; + return false; + } + if (batch_size == 0 || batch_size > kMaximumBatchSize) { + error = "StarVLA FAST batch size exceeds the runtime safety limit"; + return false; + } + if (config.time_horizon > std::numeric_limits::max() / config.action_dim) { + error = "StarVLA FAST action shape overflows size_t"; + return false; + } + per_sample = config.time_horizon * config.action_dim; + if (batch_size > std::numeric_limits::max() / per_sample) { + error = "StarVLA FAST batch shape overflows size_t"; + return false; + } + total = batch_size * per_sample; + if (total > kMaximumOutputScalars) { + error = "StarVLA FAST output tensor exceeds the runtime scalar limit"; + return false; + } + const uint64_t horizon = static_cast(config.time_horizon); + const uint64_t action_dim = static_cast(config.action_dim); + const uint64_t batch = static_cast(batch_size); + if (horizon > kMaximumIdctMultiplyAdds / horizon) { + error = "StarVLA FAST inverse DCT exceeds the runtime work limit"; + return false; + } + uint64_t multiply_adds = horizon * horizon; + if (action_dim > kMaximumIdctMultiplyAdds / multiply_adds) { + error = "StarVLA FAST inverse DCT exceeds the runtime work limit"; + return false; + } + multiply_adds *= action_dim; + if (batch > kMaximumIdctMultiplyAdds / multiply_adds) { + error = "StarVLA FAST inverse DCT exceeds the runtime work limit"; + return false; + } + return true; +} + +} // namespace + +FastCodec::FastCodec(FastCodecConfig config, std::vector> token_bytes, + std::vector fast_to_vlm_id) + : config_(config), token_bytes_(std::move(token_bytes)), + fast_to_vlm_id_(std::move(fast_to_vlm_id)) { + vlm_to_fast_id_.reserve(fast_to_vlm_id_.size()); + for (size_t fast_id = 0; fast_id < fast_to_vlm_id_.size(); ++fast_id) { + vlm_to_fast_id_.emplace_back(fast_to_vlm_id_[fast_id], static_cast(fast_id)); + } + std::sort(vlm_to_fast_id_.begin(), vlm_to_fast_id_.end()); +} + +std::unique_ptr FastCodec::create(FastCodecConfig config, + std::vector vocab_by_id, + std::vector fast_to_vlm_id, + std::string & error) { + error.clear(); + if (!std::isfinite(config.scale) || config.scale == 0.0 || config.vocab_size == 0 || + config.time_horizon == 0 || config.action_dim == 0) { + error = "StarVLA FAST codec dimensions and scale must be non-zero and finite"; + return nullptr; + } + if (config.vocab_size > static_cast(std::numeric_limits::max())) { + error = "StarVLA FAST vocabulary exceeds the int32 token-ID range"; + return nullptr; + } + if (vocab_by_id.size() != config.vocab_size || + fast_to_vlm_id.size() != config.vocab_size) { + error = "StarVLA FAST codec vocabulary or action-token map has the wrong size"; + return nullptr; + } + size_t per_sample = 0; + size_t total = 0; + if (!checked_action_count(config, 1, per_sample, total, error)) { + return nullptr; + } + std::unordered_set unique_vlm_ids; + for (int32_t vlm_id : fast_to_vlm_id) { + if (vlm_id < 0 || !unique_vlm_ids.insert(vlm_id).second) { + error = "StarVLA FAST action-token VLM IDs must be unique and non-negative"; + return nullptr; + } + } + std::vector> token_bytes; + if (!compile_token_bytes(vocab_by_id, token_bytes, error)) { + return nullptr; + } + return std::unique_ptr( + new FastCodec(config, std::move(token_bytes), std::move(fast_to_vlm_id))); +} + +std::unique_ptr FastCodec::create_compiled( + FastCodecConfig config, std::vector token_offsets, + std::vector token_bytes, std::vector fast_to_vlm_id, + std::string & error) { + error.clear(); + if (!std::isfinite(config.scale) || config.scale == 0.0 || + config.vocab_size == 0 || config.time_horizon == 0 || + config.action_dim == 0 || + config.vocab_size > + static_cast(std::numeric_limits::max())) { + error = "StarVLA FAST compiled codec dimensions and scale are invalid"; + return nullptr; + } + if (config.vocab_size == std::numeric_limits::max() || + token_offsets.size() != config.vocab_size + 1U || + fast_to_vlm_id.size() != config.vocab_size || + token_offsets.empty() || token_offsets.front() != 0 || + token_offsets.back() < 0 || + static_cast(token_offsets.back()) != token_bytes.size()) { + error = "StarVLA FAST compiled codec tensor shapes are incompatible"; + return nullptr; + } + size_t per_sample = 0; + size_t total = 0; + if (!checked_action_count(config, 1, per_sample, total, error)) { + return nullptr; + } + + std::unordered_set unique_vlm_ids; + for (int32_t vlm_id : fast_to_vlm_id) { + if (vlm_id < 0 || !unique_vlm_ids.insert(vlm_id).second) { + error = + "StarVLA FAST compiled action-token IDs must be unique and non-negative"; + return nullptr; + } + } + + std::vector> pieces; + pieces.reserve(config.vocab_size); + for (size_t index = 0; index < config.vocab_size; ++index) { + const int32_t begin = token_offsets[index]; + const int32_t end = token_offsets[index + 1U]; + if (begin < 0 || end <= begin || + static_cast(end) > token_bytes.size()) { + error = "StarVLA FAST compiled codec offsets are not strictly increasing"; + return nullptr; + } + pieces.emplace_back(token_bytes.begin() + begin, token_bytes.begin() + end); + } + return std::unique_ptr( + new FastCodec(config, std::move(pieces), std::move(fast_to_vlm_id))); +} + +std::unique_ptr FastCodec::load_hf_assets( + const std::filesystem::path & tokenizer_json, + const std::filesystem::path & processor_config_json, + const std::filesystem::path & action_token_map_json, + size_t time_horizon, size_t action_dim, std::string & error) { + error.clear(); + Json processor; + Json tokenizer; + Json action_map; + if (!read_json(processor_config_json, processor, error) || + !read_json(tokenizer_json, tokenizer, error) || + !read_json(action_token_map_json, action_map, error)) { + return nullptr; + } + + FastCodecConfig config; + std::vector vocab_by_id; + std::vector fast_to_vlm; + size_t per_sample = 0; + size_t total = 0; + if (!parse_processor_config(processor, time_horizon, action_dim, config, error) || + !checked_action_count(config, 1, per_sample, total, error) || + !parse_tokenizer_vocab(tokenizer, config.vocab_size, vocab_by_id, error) || + !parse_action_map(action_map, config.vocab_size, fast_to_vlm, error)) { + return nullptr; + } + return create(config, std::move(vocab_by_id), std::move(fast_to_vlm), error); +} + +const FastCodecConfig & FastCodec::config() const { + return config_; +} + +const std::vector & FastCodec::fast_to_vlm_ids() const { + return fast_to_vlm_id_; +} + +bool FastCodec::map_fast_to_vlm(const std::vector & fast_ids, + std::vector & vlm_ids, std::string & error) const { + vlm_ids.clear(); + error.clear(); + if (fast_ids.size() > kMaximumTokenSequence) { + error = "StarVLA FAST token sequence exceeds the runtime length limit"; + return false; + } + vlm_ids.reserve(fast_ids.size()); + for (int32_t fast_id : fast_ids) { + if (fast_id < 0 || static_cast(fast_id) >= fast_to_vlm_id_.size()) { + error = "StarVLA FAST token ID is outside the codec vocabulary"; + vlm_ids.clear(); + return false; + } + vlm_ids.push_back(fast_to_vlm_id_[static_cast(fast_id)]); + } + return true; +} + +bool FastCodec::map_vlm_to_fast(const std::vector & vlm_ids, + std::vector & fast_ids, std::string & error) const { + fast_ids.clear(); + error.clear(); + if (vlm_ids.size() > kMaximumTokenSequence) { + error = "StarVLA FAST action-token sequence exceeds the runtime length limit"; + return false; + } + fast_ids.reserve(vlm_ids.size()); + for (int32_t vlm_id : vlm_ids) { + const auto found = std::lower_bound( + vlm_to_fast_id_.begin(), vlm_to_fast_id_.end(), vlm_id, + [](const std::pair & entry, int32_t value) { + return entry.first < value; + }); + if (found == vlm_to_fast_id_.end() || found->first != vlm_id) { + error = "Qwen token ID is not present in the StarVLA FAST action-token map"; + fast_ids.clear(); + return false; + } + fast_ids.push_back(found->second); + } + return true; +} + +bool FastCodec::extract_fast_tokens(const std::vector & generated_ids, + std::vector & fast_ids, + std::string & error) const { + fast_ids.clear(); + error.clear(); + if (generated_ids.size() > kMaximumGeneratedSequence) { + error = "Qwen generated sequence exceeds the StarVLA FAST runtime length limit"; + return false; + } + for (int32_t vlm_id : generated_ids) { + const auto found = std::lower_bound( + vlm_to_fast_id_.begin(), vlm_to_fast_id_.end(), vlm_id, + [](const std::pair & entry, int32_t value) { + return entry.first < value; + }); + if (found != vlm_to_fast_id_.end() && found->first == vlm_id) { + fast_ids.push_back(found->second); + } + } + return true; +} + +bool FastCodec::byte_level_decode(const std::vector & fast_ids, + std::vector & codepoints, + std::string & error) const { + codepoints.clear(); + error.clear(); + if (fast_ids.size() > kMaximumTokenSequence) { + error = "StarVLA FAST token sequence exceeds the runtime length limit"; + return false; + } + size_t byte_count = 0; + for (int32_t fast_id : fast_ids) { + if (fast_id < 0 || static_cast(fast_id) >= token_bytes_.size()) { + error = "StarVLA FAST token ID is outside the ByteLevel BPE vocabulary"; + return false; + } + const size_t piece_size = token_bytes_[static_cast(fast_id)].size(); + if (byte_count > std::numeric_limits::max() - piece_size) { + error = "StarVLA FAST ByteLevel output size overflows size_t"; + return false; + } + byte_count += piece_size; + if (byte_count > kMaximumDecodedBytes) { + error = "StarVLA FAST ByteLevel decode exceeds the runtime byte limit"; + return false; + } + } + std::vector bytes; + bytes.reserve(byte_count); + for (int32_t fast_id : fast_ids) { + const auto & piece = token_bytes_[static_cast(fast_id)]; + bytes.insert(bytes.end(), piece.begin(), piece.end()); + } + decode_utf8_lossy(bytes, codepoints); + return true; +} + +bool FastCodec::decode_fast_tokens(const std::vector> & batch_fast_ids, + FastDecodeResult & result, std::string & error) const { + result = {}; + error.clear(); + if (batch_fast_ids.empty()) { + error = "StarVLA FAST decode batch must contain at least one sequence"; + return false; + } + if (batch_fast_ids.size() > kMaximumBatchSize) { + error = "StarVLA FAST decode batch exceeds the runtime size limit"; + return false; + } + for (const auto & fast_ids : batch_fast_ids) { + if (fast_ids.size() > kMaximumTokenSequence) { + error = "StarVLA FAST token sequence exceeds the runtime length limit"; + return false; + } + } + size_t per_sample = 0; + size_t total = 0; + if (!checked_action_count(config_, batch_fast_ids.size(), per_sample, total, error)) { + return false; + } + + result.batch_size = batch_fast_ids.size(); + result.time_horizon = config_.time_horizon; + result.action_dim = config_.action_dim; + result.actions.assign(total, 0.0); + + const double dc_scale = 1.0 / std::sqrt(static_cast(config_.time_horizon)); + const double ac_scale = std::sqrt(2.0 / static_cast(config_.time_horizon)); + for (size_t batch = 0; batch < batch_fast_ids.size(); ++batch) { + std::vector codepoints; + std::string sequence_error; + if (!byte_level_decode(batch_fast_ids[batch], codepoints, sequence_error) || + codepoints.size() != per_sample) { + error = "StarVLA FAST sequence " + std::to_string(batch) + ": " + + (sequence_error.empty() ? "decoded DCT coefficient shape mismatch" + : sequence_error); + result = {}; + return false; + } + + for (size_t action = 0; action < config_.action_dim; ++action) { + const double dc = + (static_cast(codepoints[action]) + config_.min_token) / config_.scale; + for (size_t time = 0; time < config_.time_horizon; ++time) { + double value = dc_scale * dc; + for (size_t frequency = 1; frequency < config_.time_horizon; ++frequency) { + const size_t coefficient_index = frequency * config_.action_dim + action; + const double coefficient = + (static_cast(codepoints[coefficient_index]) + config_.min_token) / + config_.scale; + const double angle = kPi * static_cast(frequency) * + static_cast(2U * time + 1U) / + (2.0 * static_cast(config_.time_horizon)); + value += ac_scale * coefficient * std::cos(angle); + } + result.actions[batch * per_sample + time * config_.action_dim + action] = value; + } + } + } + return true; +} + +bool FastCodec::decode_vlm_action_tokens( + const std::vector> & batch_vlm_ids, FastDecodeResult & result, + std::string & error) const { + if (batch_vlm_ids.empty() || batch_vlm_ids.size() > kMaximumBatchSize) { + result = {}; + error = "StarVLA FAST action-token batch is empty or exceeds the runtime size limit"; + return false; + } + std::vector> batch_fast_ids; + batch_fast_ids.reserve(batch_vlm_ids.size()); + for (const auto & vlm_ids : batch_vlm_ids) { + std::vector fast_ids; + if (!map_vlm_to_fast(vlm_ids, fast_ids, error)) { + result = {}; + return false; + } + batch_fast_ids.push_back(std::move(fast_ids)); + } + return decode_fast_tokens(batch_fast_ids, result, error); +} + +bool FastCodec::decode_generated_tokens( + const std::vector> & batch_generated_ids, FastDecodeResult & result, + std::string & error) const { + if (batch_generated_ids.empty() || batch_generated_ids.size() > kMaximumBatchSize) { + result = {}; + error = "StarVLA FAST generated-token batch is empty or exceeds the runtime size limit"; + return false; + } + std::vector> batch_fast_ids; + batch_fast_ids.reserve(batch_generated_ids.size()); + for (const auto & generated_ids : batch_generated_ids) { + std::vector fast_ids; + if (!extract_fast_tokens(generated_ids, fast_ids, error)) { + result = {}; + return false; + } + batch_fast_ids.push_back(std::move(fast_ids)); + } + return decode_fast_tokens(batch_fast_ids, result, error); +} + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/fast_codec.h b/src/models/starvla/fast_codec.h new file mode 100644 index 0000000..b84314b --- /dev/null +++ b/src/models/starvla/fast_codec.h @@ -0,0 +1,90 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +namespace robotcpp::starvla { + +struct FastCodecConfig { + double scale = 0.0; + int32_t min_token = 0; + size_t vocab_size = 0; + size_t time_horizon = 0; + size_t action_dim = 0; +}; + +struct FastDecodeResult { + size_t batch_size = 0; + size_t time_horizon = 0; + size_t action_dim = 0; + std::vector actions; +}; + +class FastCodec { +public: + static std::unique_ptr create( + FastCodecConfig config, std::vector vocab_by_id, + std::vector fast_to_vlm_id, std::string & error); + + // Constructs directly from the converter-compiled ByteLevel pieces stored + // in policy GGUF. offsets has vocab_size + 1 entries and indexes the flat + // byte buffer; no external tokenizer JSON is consulted. + static std::unique_ptr create_compiled( + FastCodecConfig config, std::vector token_offsets, + std::vector token_bytes, + std::vector fast_to_vlm_id, std::string & error); + + static std::unique_ptr load_hf_assets( + const std::filesystem::path & tokenizer_json, + const std::filesystem::path & processor_config_json, + const std::filesystem::path & action_token_map_json, + size_t time_horizon, size_t action_dim, std::string & error); + + const FastCodecConfig & config() const; + const std::vector & fast_to_vlm_ids() const; + + bool map_fast_to_vlm(const std::vector & fast_ids, + std::vector & vlm_ids, std::string & error) const; + bool map_vlm_to_fast(const std::vector & vlm_ids, + std::vector & fast_ids, std::string & error) const; + + // Extracts every mapped action token from a generated Qwen sequence in order. + // EOS stopping remains the generator's responsibility; ordinary EOS/pad/text + // IDs in the returned sequence are ignored and do not terminate this scan. + bool extract_fast_tokens(const std::vector & generated_ids, + std::vector & fast_ids, std::string & error) const; + + // Exposed for focused parity diagnostics. This is the Hugging Face ByteLevel + // decoder output before min_token adjustment and inverse DCT. + bool byte_level_decode(const std::vector & fast_ids, + std::vector & codepoints, std::string & error) const; + + bool decode_fast_tokens(const std::vector> & batch_fast_ids, + FastDecodeResult & result, std::string & error) const; + + // Strict low-level API: every input ID must be an action token. Use + // decode_generated_tokens for complete Qwen sequences containing text. + bool decode_vlm_action_tokens(const std::vector> & batch_vlm_ids, + FastDecodeResult & result, std::string & error) const; + + // Production entry point for complete Qwen generated_ids. Ordinary text and + // control tokens are filtered through the explicit inverse action-token map. + bool decode_generated_tokens(const std::vector> & batch_generated_ids, + FastDecodeResult & result, std::string & error) const; + +private: + FastCodec(FastCodecConfig config, std::vector> token_bytes, + std::vector fast_to_vlm_id); + + FastCodecConfig config_; + std::vector> token_bytes_; + std::vector fast_to_vlm_id_; + std::vector> vlm_to_fast_id_; +}; + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/fast_policy.cpp b/src/models/starvla/fast_policy.cpp new file mode 100644 index 0000000..41e2328 --- /dev/null +++ b/src/models/starvla/fast_policy.cpp @@ -0,0 +1,459 @@ +#include "models/starvla/fast_policy.h" + +#include "ggml.h" +#include "gguf.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace robotcpp::starvla { +namespace { + +constexpr const char * kArchitecture = "starvla-policy"; +constexpr const char * kActionMapTensor = + "starvla.policy.fast.action_token_map"; +constexpr const char * kOffsetsTensor = + "starvla.policy.fast.codec.token_offsets"; +constexpr const char * kTokenBytesTensor = + "starvla.policy.fast.codec.token_bytes"; + +int require_key(gguf_context * gguf, const char * key, gguf_type type) { + const int index = gguf_find_key(gguf, key); + if (index < 0) { + throw std::runtime_error(std::string("missing required FAST GGUF metadata: ") + + key); + } + if (gguf_get_kv_type(gguf, index) != type) { + throw std::runtime_error(std::string("invalid FAST GGUF metadata type: ") + + key); + } + return index; +} + +std::string require_string(gguf_context * gguf, const char * key) { + return gguf_get_val_str(gguf, require_key(gguf, key, GGUF_TYPE_STRING)); +} + +int32_t require_i32(gguf_context * gguf, const char * key) { + return gguf_get_val_i32(gguf, require_key(gguf, key, GGUF_TYPE_INT32)); +} + +float require_f32(gguf_context * gguf, const char * key) { + return gguf_get_val_f32(gguf, require_key(gguf, key, GGUF_TYPE_FLOAT32)); +} + +bool require_bool(gguf_context * gguf, const char * key) { + return gguf_get_val_bool(gguf, require_key(gguf, key, GGUF_TYPE_BOOL)); +} + +int require_array(gguf_context * gguf, const char * key, gguf_type type) { + const int index = require_key(gguf, key, GGUF_TYPE_ARRAY); + if (gguf_get_arr_type(gguf, index) != type) { + throw std::runtime_error( + std::string("invalid FAST GGUF array element type: ") + key); + } + return index; +} + +std::vector require_i32_array(gguf_context * gguf, const char * key) { + const int index = require_array(gguf, key, GGUF_TYPE_INT32); + const size_t count = gguf_get_arr_n(gguf, index); + const auto * data = + static_cast(gguf_get_arr_data(gguf, index)); + if (count != 0 && data == nullptr) { + throw std::runtime_error(std::string("missing FAST GGUF array data: ") + + key); + } + return std::vector(data, data + count); +} + +std::vector require_f32_array(gguf_context * gguf, const char * key) { + const int index = require_array(gguf, key, GGUF_TYPE_FLOAT32); + const size_t count = gguf_get_arr_n(gguf, index); + const auto * data = + static_cast(gguf_get_arr_data(gguf, index)); + if (count != 0 && data == nullptr) { + throw std::runtime_error(std::string("missing FAST GGUF array data: ") + + key); + } + return std::vector(data, data + count); +} + +std::vector require_bool_array(gguf_context * gguf, + const char * key) { + const int index = require_array(gguf, key, GGUF_TYPE_BOOL); + const size_t count = gguf_get_arr_n(gguf, index); + const auto * data = + static_cast(gguf_get_arr_data(gguf, index)); + if (count != 0 && data == nullptr) { + throw std::runtime_error(std::string("missing FAST GGUF array data: ") + + key); + } + std::vector values(count); + for (size_t i = 0; i < count; ++i) { + values[i] = data[i] != 0 ? uint8_t{1} : uint8_t{0}; + } + return values; +} + +std::vector require_string_array(gguf_context * gguf, + const char * key) { + const int index = require_array(gguf, key, GGUF_TYPE_STRING); + const size_t count = gguf_get_arr_n(gguf, index); + std::vector values; + values.reserve(count); + for (size_t i = 0; i < count; ++i) { + values.emplace_back(gguf_get_arr_str(gguf, index, i)); + } + return values; +} + +std::string profile_key(int index, const char * suffix) { + return "starvla.normalization.profile." + std::to_string(index) + "." + + suffix; +} + +struct FastRuntimeMetadata { + FastCodecConfig codec; + int token_bytes_count = 0; +}; + +FastRuntimeMetadata parse_metadata(gguf_context * gguf, + FastPolicyConfig & config) { + if (require_string(gguf, "general.architecture") != kArchitecture || + require_i32(gguf, "starvla.schema_version") != 1 || + require_string(gguf, "starvla.framework") != "fast") { + throw std::runtime_error("GGUF is not a supported StarVLA FAST policy"); + } + config.backbone_arch = require_string(gguf, "starvla.backbone.arch"); + config.bundle_uuid = require_string(gguf, "starvla.bundle.uuid"); + config.text_filename = require_string(gguf, "starvla.component.text.filename"); + config.mmproj_filename = require_string(gguf, "starvla.component.mmproj.filename"); + if (config.backbone_arch != "qwen2_5_vl" || config.bundle_uuid.empty() || + config.text_filename.empty() || config.mmproj_filename.empty()) { + throw std::runtime_error("StarVLA FAST bundle metadata is incomplete"); + } + + config.qwen_hidden_dim = require_i32(gguf, "starvla.qwen.hidden_size"); + config.qwen_input_embedding_dim = + require_i32(gguf, "starvla.qwen.input_embedding_size"); + config.qwen_vocab_size = require_i32(gguf, "starvla.qwen.vocab_size"); + config.qwen_layer_count = require_i32(gguf, "starvla.qwen.layer_count"); + config.cot_template = require_string(gguf, "starvla.prompt.cot_template"); + + config.action_dim = require_i32(gguf, "starvla.action.dimension"); + config.horizon = require_i32(gguf, "starvla.action.horizon"); + config.image_count = require_i32(gguf, "starvla.image.count"); + config.image_names = require_string_array(gguf, "starvla.image.names"); + config.image_processor_min_pixels = + require_i32(gguf, "starvla.image.processor_min_pixels"); + config.image_processor_max_pixels = + require_i32(gguf, "starvla.image.processor_max_pixels"); + config.image_patch_size = require_i32(gguf, "starvla.image.patch_size"); + config.image_spatial_merge_size = + require_i32(gguf, "starvla.image.spatial_merge_size"); + config.image_min_token_count = require_i32(gguf, "starvla.image.min_token_count"); + config.image_max_token_count = require_i32(gguf, "starvla.image.max_token_count"); + + const int max_length = require_i32(gguf, "starvla.fast.generation.max_length"); + config.generation_eos_token_ids = + require_i32_array(gguf, "starvla.fast.generation.eos_token_ids"); + config.generation_top_k = require_i32(gguf, "starvla.fast.generation.top_k"); + config.generation_repetition_penalty = + require_f32(gguf, "starvla.fast.generation.repetition_penalty"); + + FastRuntimeMetadata runtime; + runtime.codec.scale = require_f32(gguf, "starvla.fast.codec.scale"); + runtime.codec.min_token = require_i32(gguf, "starvla.fast.codec.min_token"); + runtime.codec.vocab_size = + static_cast(require_i32(gguf, "starvla.fast.codec.vocab_size")); + runtime.codec.time_horizon = + static_cast(require_i32(gguf, "starvla.fast.codec.time_horizon")); + runtime.codec.action_dim = + static_cast(require_i32(gguf, "starvla.fast.codec.action_dimension")); + const int action_token_count = + require_i32(gguf, "starvla.fast.action_token.count"); + const int offsets_count = + require_i32(gguf, "starvla.fast.codec.token_offsets_count"); + runtime.token_bytes_count = + require_i32(gguf, "starvla.fast.codec.token_bytes_count"); + + const bool valid = + config.qwen_hidden_dim > 0 && config.qwen_input_embedding_dim > 0 && + config.qwen_vocab_size > 0 && config.qwen_layer_count > 0 && + !config.cot_template.empty() && config.action_dim > 0 && config.horizon > 0 && + config.image_count > 0 && + config.image_names.size() == static_cast(config.image_count) && + config.image_processor_min_pixels > 0 && + config.image_processor_max_pixels >= config.image_processor_min_pixels && + config.image_patch_size > 0 && config.image_spatial_merge_size > 0 && + config.image_min_token_count > 0 && + config.image_max_token_count >= config.image_min_token_count && + max_length > 0 && !config.generation_eos_token_ids.empty() && + config.generation_top_k > 0 && + std::isfinite(config.generation_repetition_penalty) && + config.generation_repetition_penalty > 0.0f && + runtime.codec.vocab_size > 0 && + action_token_count == static_cast(runtime.codec.vocab_size) && + offsets_count == action_token_count + 1 && runtime.token_bytes_count > 0 && + runtime.codec.time_horizon == static_cast(config.horizon) && + runtime.codec.action_dim == static_cast(config.action_dim); + if (!valid) { + throw std::runtime_error("StarVLA FAST metadata has incompatible dimensions"); + } + config.generation_max_length = static_cast(max_length); + + NormalizationConfig & normalization = config.normalization; + normalization.clip_actions = require_bool(gguf, "starvla.normalization.clip_actions"); + normalization.binary_threshold = + require_f32(gguf, "starvla.normalization.binary_threshold"); + normalization.binary_comparison = + require_string(gguf, "starvla.normalization.binary_comparison"); + normalization.continuous_dimensions = + require_i32_array(gguf, "starvla.action.continuous_dimensions"); + normalization.binary_dimensions = + require_i32_array(gguf, "starvla.action.binary_dimensions"); + const int profile_count = require_i32(gguf, "starvla.normalization.profile_count"); + const std::vector profile_keys = + require_string_array(gguf, "starvla.normalization.profile_keys"); + if (profile_count <= 0 || profile_keys.size() != static_cast(profile_count)) { + throw std::runtime_error("StarVLA FAST normalization profiles are inconsistent"); + } + normalization.profiles.clear(); + normalization.profiles.reserve(static_cast(profile_count)); + for (int index = 0; index < profile_count; ++index) { + NormalizationProfile profile; + profile.key = require_string(gguf, profile_key(index, "key").c_str()); + profile.action_q01 = + require_f32_array(gguf, profile_key(index, "action_q01").c_str()); + profile.action_q99 = + require_f32_array(gguf, profile_key(index, "action_q99").c_str()); + profile.action_mask = + require_bool_array(gguf, profile_key(index, "action_mask").c_str()); + if (profile.key != profile_keys[static_cast(index)]) { + throw std::runtime_error("StarVLA FAST normalization profile order is inconsistent"); + } + normalization.profiles.push_back(std::move(profile)); + } + std::string normalization_error; + if (!validate_normalization_config(normalization, config.action_dim, + normalization_error)) { + throw std::runtime_error(normalization_error); + } + return runtime; +} + +struct RawTensor { + ggml_tensor * metadata = nullptr; + int index = -1; + std::vector bytes; +}; + +RawTensor read_tensor(const std::string & path, gguf_context * gguf, + ggml_context * metadata_context, const char * name, + ggml_type expected_type, int64_t expected_elements) { + RawTensor result; + result.metadata = ggml_get_tensor(metadata_context, name); + result.index = gguf_find_tensor(gguf, name); + if (result.metadata == nullptr || result.index < 0 || + result.metadata->type != expected_type || + ggml_n_dims(result.metadata) != 1 || + result.metadata->ne[0] != expected_elements || + ggml_nelements(result.metadata) != expected_elements) { + throw std::runtime_error(std::string("FAST runtime tensor shape/type mismatch: ") + + name); + } + result.bytes.resize(ggml_nbytes(result.metadata)); + std::ifstream stream(path, std::ios::binary); + if (!stream) { + throw std::runtime_error("failed to open FAST policy GGUF tensor data"); + } + const size_t offset = + gguf_get_data_offset(gguf) + + gguf_get_tensor_offset(gguf, result.index); + stream.seekg(static_cast(offset), std::ios::beg); + if (!stream || + offset > static_cast(std::numeric_limits::max())) { + throw std::runtime_error( + std::string("failed to seek FAST runtime tensor: ") + name); + } + stream.read(reinterpret_cast(result.bytes.data()), + static_cast(result.bytes.size())); + if (!stream) { + throw std::runtime_error( + std::string("failed to read FAST runtime tensor: ") + name); + } + return result; +} + +} // namespace + +struct FastPolicy::Impl { + FastPolicyConfig config; + std::unique_ptr codec; +}; + +FastPolicy::FastPolicy(std::unique_ptr impl) : impl_(std::move(impl)) {} + +FastPolicy::~FastPolicy() = default; + +std::unique_ptr FastPolicy::load(const std::string & path, + int verbosity, + std::string & error) { + error.clear(); + if (path.empty()) { + error = "StarVLA FAST policy path is required"; + return nullptr; + } + + ggml_context * metadata_context = nullptr; + gguf_init_params params{}; + params.no_alloc = true; + params.ctx = &metadata_context; + gguf_context * gguf = gguf_init_from_file(path.c_str(), params); + if (gguf == nullptr || metadata_context == nullptr) { + if (metadata_context != nullptr) { + ggml_free(metadata_context); + } + if (gguf != nullptr) { + gguf_free(gguf); + } + error = "failed to read StarVLA FAST policy GGUF"; + return nullptr; + } + auto cleanup = [&]() { + ggml_free(metadata_context); + metadata_context = nullptr; + gguf_free(gguf); + gguf = nullptr; + }; + + std::unique_ptr impl(new Impl()); + try { + const FastRuntimeMetadata runtime = parse_metadata(gguf, impl->config); + + RawTensor action_map = + read_tensor(path, gguf, metadata_context, kActionMapTensor, + GGML_TYPE_I32, static_cast(runtime.codec.vocab_size)); + RawTensor offsets = + read_tensor(path, gguf, metadata_context, kOffsetsTensor, + GGML_TYPE_I32, + static_cast(runtime.codec.vocab_size + 1)); + RawTensor token_bytes = + read_tensor(path, gguf, metadata_context, kTokenBytesTensor, + GGML_TYPE_I8, runtime.token_bytes_count); + + const uint32_t endian_probe = 1; + if (*reinterpret_cast(&endian_probe) != 1) { + throw std::runtime_error( + "FAST runtime currently requires a little-endian host"); + } + std::vector action_ids(runtime.codec.vocab_size); + std::vector token_offsets(runtime.codec.vocab_size + 1); + std::memcpy(action_ids.data(), action_map.bytes.data(), + action_map.bytes.size()); + std::memcpy(token_offsets.data(), offsets.bytes.data(), + offsets.bytes.size()); + impl->codec = FastCodec::create_compiled( + runtime.codec, std::move(token_offsets), + std::move(token_bytes.bytes), std::move(action_ids), error); + if (impl->codec == nullptr) { + throw std::runtime_error("failed to construct embedded FAST codec: " + + error); + } + if (verbosity >= 1) { + std::fprintf(stderr, + "%s: bundle=%s runtime_tensors=3 codec_vocab=%zu " + "generation_max_length=%zu profiles=%zu\n", + __func__, impl->config.bundle_uuid.c_str(), + runtime.codec.vocab_size, impl->config.generation_max_length, + impl->config.normalization.profiles.size()); + } + cleanup(); + } catch (const std::exception & exception) { + cleanup(); + error = exception.what(); + return nullptr; + } + return std::unique_ptr(new FastPolicy(std::move(impl))); +} + +bool FastPolicy::decode_generated( + const std::vector & full_sequence, + std::vector & action_token_ids, + std::vector & fast_token_ids, + std::vector & normalized_actions, + std::string & error) const { + action_token_ids.clear(); + fast_token_ids.clear(); + normalized_actions.clear(); + error.clear(); + if (impl_ == nullptr || impl_->codec == nullptr) { + error = "StarVLA FAST policy is not initialized"; + return false; + } + if (!impl_->codec->extract_fast_tokens(full_sequence, fast_token_ids, + error) || + !impl_->codec->map_fast_to_vlm(fast_token_ids, action_token_ids, + error)) { + return false; + } + FastDecodeResult decoded; + if (!impl_->codec->decode_fast_tokens({fast_token_ids}, decoded, error)) { + return false; + } + if (decoded.batch_size != 1 || + decoded.time_horizon != static_cast(impl_->config.horizon) || + decoded.action_dim != static_cast(impl_->config.action_dim) || + decoded.actions.size() != + static_cast(impl_->config.horizon * impl_->config.action_dim)) { + error = "embedded FAST codec returned an incompatible action tensor"; + return false; + } + normalized_actions.reserve(decoded.actions.size()); + for (double value : decoded.actions) { + const float converted = static_cast(value); + if (!std::isfinite(converted)) { + normalized_actions.clear(); + error = "embedded FAST codec returned a non-finite action"; + return false; + } + normalized_actions.push_back(converted); + } + return true; +} + +bool FastPolicy::unnormalize( + const std::vector & normalized_actions, + const std::string & profile_key, std::vector & actions, + std::string & error) const { + if (impl_ == nullptr) { + actions.clear(); + error = "StarVLA FAST policy is not initialized"; + return false; + } + return denormalize_actions(impl_->config.normalization, profile_key, + normalized_actions, impl_->config.horizon, + impl_->config.action_dim, actions, error); +} + +const FastPolicyConfig & FastPolicy::config() const { + if (impl_ == nullptr) { + throw std::runtime_error("StarVLA FAST policy is not initialized"); + } + return impl_->config; +} + +const char * FastPolicy::backend_name() const { + return "cpu"; +} + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/fast_policy.h b/src/models/starvla/fast_policy.h new file mode 100644 index 0000000..6a23c11 --- /dev/null +++ b/src/models/starvla/fast_policy.h @@ -0,0 +1,79 @@ +#pragma once + +#include "models/starvla/fast_codec.h" +#include "models/starvla/normalization.h" + +#include +#include +#include +#include +#include + +namespace robotcpp::starvla { + +struct FastPolicyConfig { + std::string backbone_arch; + std::string bundle_uuid; + std::string text_filename; + std::string mmproj_filename; + + int qwen_hidden_dim = 0; + int qwen_input_embedding_dim = 0; + int qwen_vocab_size = 0; + int qwen_layer_count = 0; + + std::string cot_template; + int action_dim = 0; + int horizon = 0; + + int image_count = 0; + std::vector image_names; + int image_processor_min_pixels = 0; + int image_processor_max_pixels = 0; + int image_patch_size = 0; + int image_spatial_merge_size = 0; + int image_min_token_count = 0; + int image_max_token_count = 0; + + size_t generation_max_length = 0; + std::vector generation_eos_token_ids; + int generation_top_k = 0; + float generation_repetition_penalty = 0.0f; + + NormalizationConfig normalization; +}; + +class FastPolicy { + public: + ~FastPolicy(); + + FastPolicy(const FastPolicy &) = delete; + FastPolicy & operator=(const FastPolicy &) = delete; + + static std::unique_ptr load(const std::string & path, + int verbosity, + std::string & error); + + bool decode_generated(const std::vector & full_sequence, + std::vector & action_token_ids, + std::vector & fast_token_ids, + std::vector & normalized_actions, + std::string & error) const; + + bool unnormalize(const std::vector & normalized_actions, + const std::string & profile_key, + std::vector & actions, + std::string & error) const; + + const FastPolicyConfig & config() const; + const char * backend_name() const; + + private: + struct Impl; + + explicit FastPolicy(std::unique_ptr impl); + + std::unique_ptr impl_; +}; + +} // namespace robotcpp::starvla diff --git a/tests/starvla/fast_codec_test.cpp b/tests/starvla/fast_codec_test.cpp new file mode 100644 index 0000000..701d119 --- /dev/null +++ b/tests/starvla/fast_codec_test.cpp @@ -0,0 +1,194 @@ +#include "models/starvla/fast_codec.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +using robotcpp::starvla::FastCodec; +using robotcpp::starvla::FastCodecConfig; +using robotcpp::starvla::FastDecodeResult; + +void require(bool condition, const std::string & message) { + if (!condition) { + std::cerr << "FAIL: " << message << '\n'; + std::exit(1); + } +} + +uint32_t byte_level_codepoint(uint8_t target) { + auto is_direct = [](int value) { + return (value >= 0x21 && value <= 0x7e) || (value >= 0xa1 && value <= 0xac) || + (value >= 0xae && value <= 0xff); + }; + if (is_direct(target)) { + return target; + } + uint32_t extra = 0; + for (int value = 0; value < target; ++value) { + if (!is_direct(value)) { + ++extra; + } + } + return 256U + extra; +} + +std::string utf8(uint32_t codepoint) { + std::string output; + if (codepoint <= 0x7fU) { + output.push_back(static_cast(codepoint)); + } else if (codepoint <= 0x7ffU) { + output.push_back(static_cast(0xc0U | (codepoint >> 6U))); + output.push_back(static_cast(0x80U | (codepoint & 0x3fU))); + } else { + output.push_back(static_cast(0xe0U | (codepoint >> 12U))); + output.push_back(static_cast(0x80U | ((codepoint >> 6U) & 0x3fU))); + output.push_back(static_cast(0x80U | (codepoint & 0x3fU))); + } + return output; +} + +std::unique_ptr make_synthetic_codec(size_t time_horizon = 2, + size_t action_dim = 2) { + const std::vector raw_bytes = {10, 20, 30, 40, 0xe2, 0x82, 0x28}; + std::vector vocab; + for (uint8_t byte : raw_bytes) { + vocab.push_back(utf8(byte_level_codepoint(byte))); + } + FastCodecConfig config; + config.scale = 1.0; + config.min_token = 0; + config.vocab_size = vocab.size(); + config.time_horizon = time_horizon; + config.action_dim = action_dim; + std::string error; + auto codec = FastCodec::create(config, vocab, {100, 42, 999, 7, 501, 502, 503}, error); + require(codec != nullptr, "synthetic FAST codec must construct: " + error); + return codec; +} + +void run_unit_tests() { + auto codec = make_synthetic_codec(); + std::string error; + + std::vector vlm_ids; + require(codec->map_fast_to_vlm({3, 0, 2, 1}, vlm_ids, error), + "FAST-to-VLM mapping must succeed"); + require(vlm_ids == std::vector({7, 100, 999, 42}), + "FAST-to-VLM mapping must use the explicit non-contiguous table"); + + std::vector fast_ids; + require(codec->map_vlm_to_fast(vlm_ids, fast_ids, error), + "VLM-to-FAST mapping must succeed"); + require(fast_ids == std::vector({3, 0, 2, 1}), + "VLM-to-FAST mapping must invert the explicit table"); + require(!codec->map_vlm_to_fast({101}, fast_ids, error), + "an unmapped Qwen token must be rejected"); + + require(codec->extract_fast_tokens({-1, 42, 1234, 100, 7, 42}, fast_ids, error), + "full Qwen sequence action extraction must succeed"); + require(fast_ids == std::vector({1, 0, 3, 1}), + "action extraction must filter with the inverse map and preserve order"); + + std::vector codepoints; + require(codec->byte_level_decode({0, 1, 2, 3}, codepoints, error), + "ByteLevel decode must succeed"); + require(codepoints == std::vector({10, 20, 30, 40}), + "ByteLevel decode must invert the GPT-2 byte alphabet"); + require(codec->byte_level_decode({4, 5, 6}, codepoints, error), + "lossy UTF-8 ByteLevel decode must succeed"); + require(codepoints == std::vector({0xfffdU, 0x28U}), + "ByteLevel decode must match Rust UTF-8 replacement semantics"); + + FastDecodeResult decoded; + require(codec->decode_fast_tokens({{0, 1, 2, 3}}, decoded, error), + "synthetic inverse DCT must succeed"); + require(decoded.actions.size() == 4, + "valid synthetic tokens must produce one 2x2 action chunk"); + const double root_half = std::sqrt(0.5); + const std::vector expected = { + root_half * (10.0 + 30.0), root_half * (20.0 + 40.0), + root_half * (10.0 - 30.0), root_half * (20.0 - 40.0), + }; + for (size_t index = 0; index < expected.size(); ++index) { + require(std::abs(decoded.actions[index] - expected[index]) < 1e-12, + "orthonormal inverse DCT must match the analytical result"); + } + + require(!codec->decode_fast_tokens({{}, {0, 1, 2}, {9999}}, decoded, error), + "malformed FAST coefficients must fail"); + require(decoded.actions.empty(), "failed FAST decode must not return zero actions"); + require(!codec->decode_fast_tokens({}, decoded, error), "an empty batch must fail explicitly"); + + require(!codec->decode_fast_tokens( + {{0, 1, 2, 3}, {9999}, {0, 1, 2, 3}}, decoded, error), + "a malformed FAST batch member must fail the batch"); + + FastDecodeResult decoded_generated; + require(codec->decode_generated_tokens( + {{-1, 100, 123456, 42, 999, 555555, 7}}, decoded_generated, error), + "complete generated_ids must filter then decode"); + bool generated_matches_expected = decoded_generated.actions.size() == expected.size(); + for (size_t index = 0; generated_matches_expected && index < expected.size(); ++index) { + generated_matches_expected = + std::abs(decoded_generated.actions[index] - expected[index]) < 1e-12; + } + require(generated_matches_expected, + "complete generated_ids must match the pure FAST-token action decode"); + require(!codec->decode_generated_tokens({{1, 2, 3}}, decoded_generated, error), + "a generated sequence without action tokens must fail"); + require(!codec->decode_vlm_action_tokens({{123456}}, decoded_generated, error), + "strict low-level VLM action-token decode must reject ordinary Qwen tokens"); + + std::vector maximum_generated_sequence(2048, 123456); + require(!codec->decode_generated_tokens( + {maximum_generated_sequence}, decoded_generated, error), + "max_length text without action tokens must fail decode"); + std::vector oversized_sequence(2049, 100); + require(!codec->extract_fast_tokens(oversized_sequence, fast_ids, error), + "generated token sequences beyond official max_length must fail before allocation"); + + FastCodecConfig invalid_config = codec->config(); + invalid_config.scale = 0.0; + auto invalid = FastCodec::create(invalid_config, + {"a", "b", "c", "d", "e", "f", "g"}, + {0, 1, 2, 3, 4, 5, 6}, error); + require(invalid == nullptr, "zero FAST scale must be rejected"); + + invalid_config = codec->config(); + invalid = FastCodec::create(invalid_config, + {"a", "b", "c", "d", "e", "f", "g"}, + {0, 1, 2, 3, 4, 5, 5}, error); + require(invalid == nullptr, "duplicate VLM action-token IDs must be rejected"); + + invalid_config = codec->config(); + invalid_config.time_horizon = 1025; + invalid = FastCodec::create(invalid_config, + {"a", "b", "c", "d", "e", "f", "g"}, + {0, 1, 2, 3, 4, 5, 6}, error); + require(invalid == nullptr, "oversized FAST horizons must fail before graph work"); + + std::vector> oversized_batch(1025); + require(!codec->decode_fast_tokens(oversized_batch, decoded, error), + "oversized FAST batches must fail before output allocation"); + + auto work_limited_codec = make_synthetic_codec(257, 1); + std::vector> excessive_idct_batch(1024); + require(!work_limited_codec->decode_fast_tokens(excessive_idct_batch, decoded, error) && + error.find("work limit") != std::string::npos, + "inverse-DCT work accounting must include the full batch dimension"); +} + +} // namespace + +int main() { + run_unit_tests(); + std::cout << "starvla FAST codec unit tests passed\n"; + return 0; +} diff --git a/tests/starvla/fast_runtime_test.cpp b/tests/starvla/fast_runtime_test.cpp new file mode 100644 index 0000000..7ec1dc7 --- /dev/null +++ b/tests/starvla/fast_runtime_test.cpp @@ -0,0 +1,137 @@ +#include "models/starvla/fast_codec.h" +#include "models/starvla/fast_policy.h" +#include "models/starvla/qwen3vl_bridge.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +using robotcpp::starvla::FastCodec; +using robotcpp::starvla::FastCodecConfig; +using robotcpp::starvla::FastDecodeResult; +using robotcpp::starvla::FastPolicy; + +void require(bool condition, const std::string & message) { + if (!condition) { + std::cerr << "FAIL: " << message << '\n'; + std::exit(1); + } +} + +void test_generation_selector() { + std::string error; + int32_t token = -1; + std::vector logits = {10.0f, 9.0f, 1.0f}; + require(robotcpp::starvla::qwen_vl_select_repetition_penalized_top1( + logits.data(), logits.size(), {0}, 2.0f, token, error), + "repetition-penalized selector must succeed: " + error); + require(token == 1, + "a repeated positive logit must be divided before top_k=1"); + + logits = {-1.0f, -1.5f, -4.0f}; + require(robotcpp::starvla::qwen_vl_select_repetition_penalized_top1( + logits.data(), logits.size(), {0}, 2.0f, token, error), + "negative-logit selector must succeed"); + require(token == 1, + "a repeated negative logit must be multiplied before top_k=1"); + + logits = {3.0f, 3.0f}; + require(robotcpp::starvla::qwen_vl_select_repetition_penalized_top1( + logits.data(), logits.size(), {}, 1.05f, token, error) && + token == 0, + "top_k=1 tie handling must match torch.argmax first-index semantics"); + require(!robotcpp::starvla::qwen_vl_select_repetition_penalized_top1( + logits.data(), logits.size(), {2}, 1.05f, token, error), + "out-of-vocabulary history must be rejected"); + logits[0] = std::numeric_limits::quiet_NaN(); + require(!robotcpp::starvla::qwen_vl_select_repetition_penalized_top1( + logits.data(), logits.size(), {}, 1.05f, token, error), + "NaN generation logits must fail closed"); +} + +void test_compiled_codec() { + FastCodecConfig config; + config.scale = 1.0; + config.min_token = 0; + config.vocab_size = 4; + config.time_horizon = 2; + config.action_dim = 2; + std::string error; + auto codec = FastCodec::create_compiled( + config, {0, 1, 2, 3, 4}, {10, 20, 30, 40}, + {100, 101, 102, 103}, error); + require(codec != nullptr, + "compiled FAST codec must construct without sidecars: " + error); + + FastDecodeResult decoded; + require(codec->decode_generated_tokens( + {{999, 100, 101, 888, 102, 103}}, decoded, error), + "compiled FAST codec must filter and decode a full Qwen sequence"); + const double root_half = std::sqrt(0.5); + const std::vector expected = { + root_half * 40.0, root_half * 60.0, + root_half * -20.0, root_half * -20.0, + }; + require(decoded.actions.size() == expected.size(), + "compiled FAST codec must return the configured 2x2 action shape"); + for (size_t i = 0; i < expected.size(); ++i) { + require(std::fabs(decoded.actions[i] - expected[i]) < 1.0e-12, + "compiled FAST codec IDCT differs from the analytical result"); + } + + require(FastCodec::create_compiled( + config, {0, 1, 1, 3, 4}, {10, 20, 30, 40}, + {100, 101, 102, 103}, error) == nullptr, + "compiled FAST codec must reject non-increasing offsets"); +} + +void test_policy(const std::string & path) { + std::string error; + std::unique_ptr policy = FastPolicy::load(path, 0, error); + require(policy != nullptr, "official FAST policy GGUF must load: " + error); + const auto & config = policy->config(); + require(config.bundle_uuid == "b2651406-918b-524b-9df6-66861d744f29" && + config.qwen_hidden_dim == 2048 && + config.qwen_vocab_size == 153713 && + config.generation_max_length == 2048 && + config.generation_eos_token_ids == + std::vector({151645, 151643}), + "official FAST policy metadata must expose the pinned runtime"); + + std::vector action_ids; + std::vector fast_ids; + std::vector normalized; + require(!policy->decode_generated( + {100, 151665, 200}, action_ids, fast_ids, normalized, error), + "incomplete FAST output must fail"); + require(action_ids == std::vector({151665}) && + fast_ids == std::vector({0}) && normalized.empty(), + "failed FAST decode must not return actions"); +} + +} // namespace + +int main(int argc, char ** argv) { + test_generation_selector(); + test_compiled_codec(); + + if (argc == 3 && std::string(argv[1]) == "--policy") { + test_policy(argv[2]); + } else if (argc == 3 && std::string(argv[1]) == "--expect-reject") { + std::string error; + require(FastPolicy::load(argv[2], 0, error) == nullptr && !error.empty(), + "tampered FAST policy GGUF must fail closed"); + } else if (argc != 1) { + std::cerr << "usage: " << argv[0] + << " [--policy|--expect-reject ]\n"; + return 2; + } + return 0; +} diff --git a/tests/starvla/test_starvla_qwen25_fast.py b/tests/starvla/test_starvla_qwen25_fast.py new file mode 100644 index 0000000..a8c7486 --- /dev/null +++ b/tests/starvla/test_starvla_qwen25_fast.py @@ -0,0 +1,111 @@ +from __future__ import annotations + +import sys +import tempfile +import unittest +from pathlib import Path + +import numpy as np + + +REPO_ROOT = Path(__file__).resolve().parents[2] +TOOLS_DIR = REPO_ROOT / "tools" / "hf2gguf" / "starvla" +sys.path.insert(0, str(TOOLS_DIR)) + +import convert_starvla_qwen25_fast as converter # noqa: E402 +from starvla_checkpoint import ( # noqa: E402 + StarVLAError, + load_catalog, + official_bundle_uuid, +) + + +SOURCE_ROOT = REPO_ROOT / "ckpts" / "starvla" / "sources" +POLICY_DIR = SOURCE_ROOT / "qwen25-fast-bridge-rt1" / "d9e2977d21755e78a0dd5f9a61586075a636d669" +QWEN_DIR = SOURCE_ROOT / "qwen2.5-vl-3b-instruct-action" / "ce86bd9a53416527b8361e8dfc47316288ffa110" +CODEC_DIR = SOURCE_ROOT / "fast-codec" / "ec4d7aa71691cac0b8bed6942be45684db2110f4" + + +def runtime_inputs() -> tuple[dict[str, object], dict[str, object], dict[str, object]]: + catalog = load_catalog() + entry, qwen, codec = converter.validate_catalog_contract(catalog) + manifest = { + "bundle_uuid": official_bundle_uuid(entry, catalog), + "source": { + "starvla_revision": catalog["source_revisions"]["starvla"], + "llama_cpp_revision": catalog["source_revisions"]["llama_cpp"], + "qwen_repo_id": qwen["repo_id"], + "qwen_revision": qwen["revision"], + }, + } + return manifest, entry, codec + + +@unittest.skipUnless( + POLICY_DIR.is_dir() and QWEN_DIR.is_dir() and CODEC_DIR.is_dir(), + "pinned FAST assets are not available", +) +class Qwen25FastTest(unittest.TestCase): + def test_preflight_and_codec_tables(self) -> None: + report = converter.preflight(load_catalog(), POLICY_DIR, QWEN_DIR, CODEC_DIR) + arrays = converter.compile_fast_runtime_tensors(QWEN_DIR, CODEC_DIR) + + self.assertEqual(report["variant"], "qwen25_fast") + self.assertEqual(set(arrays), converter.FAST_RUNTIME_TENSOR_NAMES) + self.assertEqual( + arrays[converter.ACTION_TOKEN_MAP_TENSOR].shape, + (converter.ACTION_TOKEN_COUNT,), + ) + self.assertEqual( + arrays[converter.CODEC_TOKEN_OFFSETS_TENSOR].shape, + (converter.ACTION_TOKEN_COUNT + 1,), + ) + + def test_runtime_policy_round_trip(self) -> None: + manifest, entry, codec = runtime_inputs() + metadata, arrays = converter.build_fast_runtime_policy( + manifest=manifest, + entry=entry, + codec_entry=codec, + source_dir=POLICY_DIR, + qwen_dir=QWEN_DIR, + codec_dir=CODEC_DIR, + ) + self.assertEqual(metadata["starvla.framework"], "fast") + self.assertEqual(metadata["starvla.model_type"], "starvla") + self.assertNotIn("starvla.fast.codec.decode_fallback", metadata) + self.assertNotIn("starvla.fast.runtime_contract_json", metadata) + + with tempfile.TemporaryDirectory() as temporary: + path = Path(temporary) / "policy.gguf" + converter.write_fast_runtime_policy_gguf(path, metadata, arrays) + record = converter.validate_fast_runtime_policy_gguf( + path, + expected_metadata=metadata, + expected_arrays=arrays, + ) + self.assertEqual(record["tensor_count"], 3) + + def test_validator_rejects_duplicate_action_token_ids(self) -> None: + manifest, entry, codec = runtime_inputs() + metadata, arrays = converter.build_fast_runtime_policy( + manifest=manifest, + entry=entry, + codec_entry=codec, + source_dir=POLICY_DIR, + qwen_dir=QWEN_DIR, + codec_dir=CODEC_DIR, + ) + arrays = dict(arrays) + action_map = np.array(arrays[converter.ACTION_TOKEN_MAP_TENSOR], copy=True) + action_map[1] = action_map[0] + arrays[converter.ACTION_TOKEN_MAP_TENSOR] = action_map + with tempfile.TemporaryDirectory() as temporary: + path = Path(temporary) / "policy.gguf" + converter.write_fast_runtime_policy_gguf(path, metadata, arrays) + with self.assertRaisesRegex(StarVLAError, "codec tensors"): + converter.validate_fast_runtime_policy_gguf(path) + + +if __name__ == "__main__": + unittest.main() diff --git a/tools/hf2gguf/starvla/convert_starvla_qwen25_fast.py b/tools/hf2gguf/starvla/convert_starvla_qwen25_fast.py new file mode 100644 index 0000000..8c17940 --- /dev/null +++ b/tools/hf2gguf/starvla/convert_starvla_qwen25_fast.py @@ -0,0 +1,1542 @@ +#!/usr/bin/env python3 +"""Stage and convert the official Qwen2.5-VL StarVLA FAST checkpoint.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +import shutil +import subprocess +import sys +from pathlib import Path +from typing import Any, Mapping, Sequence + +import numpy as np + +from convert_starvla_qwen_to_gguf import build_commands, verify_llama_checkout +from starvla_checkpoint import ( + DEFAULT_CATALOG, + StarVLAError, + atomic_write_json, + build_inventory, + get_qwen_asset, + get_variant, + inventory_summary, + load_catalog, + load_checkpoint_state, + official_bundle_uuid, + sha256_file, + staged_qwen_asset_hashes, + validate_qwen_vlm_destination_names, + verify_catalog_files, + verify_checkpoint_file, + verify_staged_assets, + verify_staged_shards, +) +from starvla_surgery import ( + copy_policy_assets, + copy_qwen_assets, + parse_size, + write_safetensor_shards, +) + + +REPOSITORY_ROOT = Path(__file__).resolve().parents[3] +LLAMA_GGUF_PY = REPOSITORY_ROOT / "third_party" / "llama.cpp" / "gguf-py" +if not LLAMA_GGUF_PY.is_dir(): + raise ImportError( + "third_party/llama.cpp/gguf-py is required; initialize the llama.cpp submodule" + ) +sys.path.insert(0, str(LLAMA_GGUF_PY)) + +import gguf # noqa: E402 + + +VARIANT_KEY = "qwen25_fast" +QWEN_ASSET_KEY = "qwen2_5_vl_3b_instruct_action" +FAST_CODEC_ASSET_KEY = "fast_codec" + +MODEL_TYPE = "starvla" +BACKBONE = "qwen2_5_vl" +FRAMEWORK = "fast" + +ACTION_TOKEN_COUNT = 2048 +ACTION_TOKEN_MIN = 151665 +ACTION_TOKEN_MAX = 153712 +ACTION_DIM = 7 +ACTION_HORIZON = 16 +MAX_LENGTH = 2048 + +BOS_TOKEN_ID = 151643 +EOS_TOKEN_IDS = [151645, 151643] +PAD_TOKEN_ID = 151643 +GENERATION_CONTRACT = { + "max_length": MAX_LENGTH, + "do_sample": True, + "temperature": 0.1, + "top_k": 1, + "top_p": 0.001, + "repetition_penalty": 1.05, + "bos_token_id": BOS_TOKEN_ID, + "eos_token_id": EOS_TOKEN_IDS, + "pad_token_id": PAD_TOKEN_ID, +} + +EXPECTED_INVENTORY = { + "total_tensors": 825, + "vlm_tensors": 825, + "policy_tensors": 0, + "visual_tensors": 390, + "text_tensors": 434, + "lm_head_tensors": 1, + "total_numel": 4_073_066_496, + "vlm_numel": 4_073_066_496, + "policy_numel": 0, + "total_nbytes": 8_146_132_992, + "vlm_nbytes": 8_146_132_992, + "policy_nbytes": 0, + "dtypes": {"bfloat16": 825}, + "storage_alias_groups": 0, +} + +TEXT_FILENAME = "qwen-qwen25-fast-bf16.gguf" +MMPROJ_FILENAME = "mmproj-qwen25-fast-bf16.gguf" +POLICY_FILENAME = "policy-qwen25-fast.gguf" +STAGING_MANIFEST_FILENAME = "qwen25-fast-staging-manifest.json" +BUNDLE_MANIFEST_FILENAME = "qwen25-fast-bundle-manifest.json" + +COT_PROMPT = ( + "Your task is {instruction}. To identify the key objects for your task. " + "Locate their bounding boxes in [x1,y1,x2,y2] format." +) +ACTION_NAMES = ["x", "y", "z", "roll", "pitch", "yaw", "gripper"] +EXPECTED_NORMALIZATION_PROFILES = { + "bridge_dataset", + "fractal20220817_data", +} + +ACTION_TOKEN_MAP_TENSOR = "starvla.policy.fast.action_token_map" +CODEC_TOKEN_OFFSETS_TENSOR = "starvla.policy.fast.codec.token_offsets" +CODEC_TOKEN_BYTES_TENSOR = "starvla.policy.fast.codec.token_bytes" +FAST_RUNTIME_TENSOR_NAMES = { + ACTION_TOKEN_MAP_TENSOR, + CODEC_TOKEN_OFFSETS_TENSOR, + CODEC_TOKEN_BYTES_TENSOR, +} +QWEN25VL_PROCESSOR_MIN_PIXELS = 3_136 +QWEN25VL_PROCESSOR_MAX_PIXELS = 12_845_056 +QWEN25VL_IMAGE_PATCH_SIZE = 14 +QWEN25VL_TEMPORAL_PATCH_SIZE = 2 +QWEN25VL_SPATIAL_MERGE_SIZE = 2 +QWEN25VL_MIN_IMAGE_TOKENS = 4 +QWEN25VL_MAX_IMAGE_TOKENS = 16_384 +QWEN25VL_IMAGE_MEAN = [0.48145466, 0.4578275, 0.40821073] +QWEN25VL_IMAGE_STD = [0.26862954, 0.26130258, 0.27577711] + +def load_json_object(path: Path) -> dict[str, Any]: + try: + value = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise StarVLAError(f"failed to load JSON {path}: {exc}") from exc + if not isinstance(value, dict): + raise StarVLAError(f"expected a JSON object in {path}") + return value + + +def validate_catalog_contract( + catalog: Mapping[str, Any], +) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]: + entry = get_variant(catalog, VARIANT_KEY) + expected = { + "framework": FRAMEWORK, + "backbone": BACKBONE, + "model_type": MODEL_TYPE, + "status": "official_policy", + "qwen_asset": QWEN_ASSET_KEY, + "policy_prefixes": [], + } + mismatches = [ + f"{key}: expected {value!r}, got {entry.get(key)!r}" + for key, value in expected.items() + if entry.get(key) != value + ] + checkpoint = entry.get("checkpoint") + if ( + not isinstance(checkpoint, Mapping) + or checkpoint.get("path") != "checkpoints/steps_10000_pytorch_model.pt" + or checkpoint.get("size") != 8_146_439_050 + or checkpoint.get("sha256") + != "f30e89a6b2a166fa3f48af42d5cffde07be44074b861abc7b57e1ccdb734e81e" + ): + mismatches.append("checkpoint: not the reviewed steps_10000 source lock") + if entry.get("policy_tensors") not in (None, []): + mismatches.append("policy_tensors: FAST must not split a separate policy head") + official_bundle_uuid(entry, catalog) + + qwen_name, qwen_entry = get_qwen_asset(catalog, entry) + if qwen_name != QWEN_ASSET_KEY: + mismatches.append(f"Qwen asset: expected {QWEN_ASSET_KEY!r}, got {qwen_name!r}") + codec_entry = catalog.get("shared_assets", {}).get(FAST_CODEC_ASSET_KEY) + if not isinstance(codec_entry, dict): + mismatches.append("FAST codec: missing pinned shared asset") + codec_entry = {} + if mismatches: + raise StarVLAError("Qwen2.5 FAST catalog contract mismatch: " + "; ".join(mismatches)) + return entry, qwen_entry, codec_entry + + +def validate_qwen_config(qwen_dir: Path) -> dict[str, Any]: + config = load_json_object(qwen_dir / "config.json") + text = config.get("text_config") + vision = config.get("vision_config") + if not isinstance(text, dict) or not isinstance(vision, dict): + raise StarVLAError("Qwen2.5 FAST config has no text_config/vision_config object") + expected_top = { + "architectures": ["Qwen2_5_VLForConditionalGeneration"], + "model_type": "qwen2_5_vl", + "dtype": "bfloat16", + "vocab_size": 151936, + "image_token_id": 151655, + "video_token_id": 151656, + "vision_token_id": 151654, + "vision_start_token_id": 151652, + "vision_end_token_id": 151653, + } + expected_text = { + "model_type": "qwen2_5_vl_text", + "dtype": "bfloat16", + "hidden_size": 2048, + "intermediate_size": 11008, + "num_hidden_layers": 36, + "num_attention_heads": 16, + "num_key_value_heads": 2, + "vocab_size": 153713, + "tie_word_embeddings": True, + } + expected_vision = { + "depth": 32, + "hidden_size": 1280, + "intermediate_size": 3420, + "num_heads": 16, + "out_hidden_size": 2048, + "patch_size": 14, + "spatial_merge_size": 2, + "temporal_patch_size": 2, + "window_size": 112, + "fullatt_block_indexes": [7, 15, 23, 31], + } + mismatches = [] + for owner_name, owner, contract in ( + ("config", config, expected_top), + ("text_config", text, expected_text), + ("vision_config", vision, expected_vision), + ): + for key, value in contract.items(): + if owner.get(key) != value: + mismatches.append( + f"{owner_name}.{key}: expected {value!r}, got {owner.get(key)!r}" + ) + if mismatches: + raise StarVLAError("Qwen2.5 FAST config mismatch: " + "; ".join(mismatches)) + return { + "text_hidden_size": 2048, + "text_layers": 36, + "text_attention_heads": 16, + "text_key_value_heads": 2, + "text_attention_head_dim": 128, + "vision_hidden_size": 1280, + "vision_layers": 32, + "vision_attention_heads": 16, + "vision_full_attention_blocks": [7, 15, 23, 31], + "vision_deepstack": False, + "vocab_size": 153713, + } + + +def validate_qwen_processor(qwen_dir: Path) -> dict[str, Any]: + config = load_json_object(qwen_dir / "preprocessor_config.json") + expected = { + "do_convert_rgb": True, + "do_normalize": True, + "do_rescale": True, + "do_resize": True, + "image_mean": QWEN25VL_IMAGE_MEAN, + "image_std": QWEN25VL_IMAGE_STD, + "image_processor_type": "Qwen2VLImageProcessorFast", + "max_pixels": QWEN25VL_PROCESSOR_MAX_PIXELS, + "merge_size": QWEN25VL_SPATIAL_MERGE_SIZE, + "min_pixels": QWEN25VL_PROCESSOR_MIN_PIXELS, + "patch_size": QWEN25VL_IMAGE_PATCH_SIZE, + "processor_class": "Qwen2_5_VLProcessor", + "resample": 3, + "rescale_factor": 1.0 / 255.0, + "temporal_patch_size": QWEN25VL_TEMPORAL_PATCH_SIZE, + } + mismatches = [ + f"{key}: expected {value!r}, got {config.get(key)!r}" + for key, value in expected.items() + if config.get(key) != value + ] + expected_size = { + "longest_edge": QWEN25VL_PROCESSOR_MAX_PIXELS, + "shortest_edge": QWEN25VL_PROCESSOR_MIN_PIXELS, + } + if config.get("size") != expected_size: + mismatches.append( + f"size: expected {expected_size!r}, got {config.get('size')!r}" + ) + chat_template = qwen_dir / "chat_template.jinja" + if not chat_template.is_file() or chat_template.stat().st_size == 0: + mismatches.append("chat_template.jinja: missing or empty") + if mismatches: + raise StarVLAError( + "Qwen2.5 FAST processor contract mismatch: " + "; ".join(mismatches) + ) + return { + "min_pixels": QWEN25VL_PROCESSOR_MIN_PIXELS, + "max_pixels": QWEN25VL_PROCESSOR_MAX_PIXELS, + "patch_size": QWEN25VL_IMAGE_PATCH_SIZE, + "temporal_patch_size": QWEN25VL_TEMPORAL_PATCH_SIZE, + "merge_size": QWEN25VL_SPATIAL_MERGE_SIZE, + "min_image_tokens": QWEN25VL_MIN_IMAGE_TOKENS, + "max_image_tokens": QWEN25VL_MAX_IMAGE_TOKENS, + "image_mean": list(QWEN25VL_IMAGE_MEAN), + "image_std": list(QWEN25VL_IMAGE_STD), + "chat_template_sha256": sha256_file(chat_template), + } + + +def validate_generation_config(qwen_dir: Path) -> dict[str, Any]: + config = load_json_object(qwen_dir / "generation_config.json") + mismatches = [ + f"{key}: expected {value!r}, got {config.get(key)!r}" + for key, value in GENERATION_CONTRACT.items() + if key != "max_length" and config.get(key) != value + ] + if mismatches: + raise StarVLAError( + "Qwen2.5 FAST generation_config mismatch: " + "; ".join(mismatches) + ) + return dict(GENERATION_CONTRACT) + + +def _action_mapping(path: Path) -> dict[str, int]: + raw = load_json_object(path) + mapping: dict[str, int] = {} + for key, value in raw.items(): + if not isinstance(value, int) or isinstance(value, bool): + raise StarVLAError(f"FAST action mapping has a non-integer ID for {key!r}") + mapping[key] = value + return mapping + + +def validate_action_token_mapping(qwen_dir: Path) -> dict[str, Any]: + expected = { + f"": ACTION_TOKEN_MIN + index + for index in range(ACTION_TOKEN_COUNT) + } + primary = _action_mapping(qwen_dir / "added_token_id_map.json") + added = { + key: value + for key, value in _action_mapping(qwen_dir / "added_tokens.json").items() + if key.startswith(" Qwen 151665..153712 mapping" + ) + if added != expected: + raise StarVLAError("added_tokens.json disagrees with the pinned FAST action mapping") + + tokenizer = load_json_object(qwen_dir / "tokenizer.json") + tokenizer_action = { + str(record.get("content")): record + for record in tokenizer.get("added_tokens", []) + if isinstance(record, dict) + and str(record.get("content", "")).startswith(" dict[str, Any]: + hashes = verify_catalog_files(codec_dir, codec_entry) + config = load_json_object(codec_dir / "processor_config.json") + expected = { + "processor_class": "UniversalActionProcessor", + "scale": 10, + "vocab_size": ACTION_TOKEN_COUNT, + "min_token": -354, + "action_dim": None, + "time_horizon": None, + } + mismatches = [ + f"{key}: expected {value!r}, got {config.get(key)!r}" + for key, value in expected.items() + if config.get(key) != value + ] + if mismatches: + raise StarVLAError("FAST codec config mismatch: " + "; ".join(mismatches)) + return { + "repo_id": codec_entry["repo_id"], + "revision": codec_entry["revision"], + "scale": 10, + "min_token": -354, + "vocab_size": ACTION_TOKEN_COUNT, + "action_dim": ACTION_DIM, + "time_horizon": ACTION_HORIZON, + "files": hashes, + } + + +def _sha256_bytes(value: bytes) -> str: + return hashlib.sha256(value).hexdigest() + + +def _byte_level_inverse_alphabet() -> dict[int, int]: + direct = { + *range(0x21, 0x7F), + *range(0xA1, 0xAD), + *range(0xAE, 0x100), + } + inverse = {value: value for value in direct} + extra = 0 + for value in range(0x100): + if value not in direct: + inverse[0x100 + extra] = value + extra += 1 + if extra != 68 or len(inverse) != 256: + raise AssertionError("internal GPT-2 ByteLevel alphabet construction drift") + return inverse + + +def compile_fast_runtime_tensors( + qwen_dir: Path, + codec_dir: Path, +) -> dict[str, np.ndarray]: + """Compile the pinned HF FAST decode assets into runtime-only integer tables.""" + processor = load_json_object(codec_dir / "processor_config.json") + expected_processor = { + "processor_class": "UniversalActionProcessor", + "scale": 10, + "vocab_size": ACTION_TOKEN_COUNT, + "min_token": -354, + "action_dim": None, + "time_horizon": None, + } + mismatches = [ + f"{key}: expected {value!r}, got {processor.get(key)!r}" + for key, value in expected_processor.items() + if processor.get(key) != value + ] + if mismatches: + raise StarVLAError( + "FAST runtime processor contract mismatch: " + "; ".join(mismatches) + ) + + tokenizer = load_json_object(codec_dir / "tokenizer.json") + decoder = tokenizer.get("decoder") + model = tokenizer.get("model") + if ( + tokenizer.get("version") != "1.0" + or tokenizer.get("added_tokens") != [] + or decoder + != { + "type": "ByteLevel", + "add_prefix_space": True, + "trim_offsets": True, + "use_regex": True, + } + or not isinstance(model, dict) + or model.get("type") != "BPE" + or not isinstance(model.get("vocab"), dict) + ): + raise StarVLAError( + "FAST tokenizer is not the pinned ByteLevel BPE decode contract" + ) + vocab = model["vocab"] + if len(vocab) != ACTION_TOKEN_COUNT: + raise StarVLAError( + f"FAST tokenizer vocabulary must contain {ACTION_TOKEN_COUNT} entries" + ) + vocab_by_id: list[str | None] = [None] * ACTION_TOKEN_COUNT + for piece, token_id in vocab.items(): + if ( + not isinstance(piece, str) + or not piece + or isinstance(token_id, bool) + or not isinstance(token_id, int) + or token_id < 0 + or token_id >= ACTION_TOKEN_COUNT + or vocab_by_id[token_id] is not None + ): + raise StarVLAError("FAST tokenizer vocabulary IDs are not an exact bijection") + vocab_by_id[token_id] = piece + if any(piece is None for piece in vocab_by_id): + raise StarVLAError("FAST tokenizer vocabulary has missing IDs") + + inverse_alphabet = _byte_level_inverse_alphabet() + token_offsets = [0] + flattened = bytearray() + for token_id, optional_piece in enumerate(vocab_by_id): + if optional_piece is None: + raise AssertionError("FAST vocabulary completeness check failed") + for character in optional_piece: + byte_value = inverse_alphabet.get(ord(character)) + if byte_value is None: + raise StarVLAError( + "FAST tokenizer piece contains a code point outside the " + f"ByteLevel alphabet at token ID {token_id}" + ) + flattened.append(byte_value) + token_offsets.append(len(flattened)) + + validate_action_token_mapping(qwen_dir) + raw_mapping = _action_mapping(qwen_dir / "added_token_id_map.json") + action_token_map = np.asarray( + [ + raw_mapping[f""] + for token_id in range(ACTION_TOKEN_COUNT) + ], + dtype=np.int32, + ) + offsets = np.asarray(token_offsets, dtype=np.int32) + token_bytes = ( + np.frombuffer(bytes(flattened), dtype=np.uint8).view(np.int8).copy() + ) + + if ( + offsets.shape != (ACTION_TOKEN_COUNT + 1,) + or offsets[0] != 0 + or offsets[-1] != token_bytes.size + or np.any(np.diff(offsets) <= 0) + or action_token_map.tolist() + != list(range(ACTION_TOKEN_MIN, ACTION_TOKEN_MAX + 1)) + ): + raise StarVLAError("compiled FAST runtime tensor shape/content mismatch") + return { + ACTION_TOKEN_MAP_TENSOR: action_token_map, + CODEC_TOKEN_OFFSETS_TENSOR: offsets, + CODEC_TOKEN_BYTES_TENSOR: token_bytes, + } + + +def normalization_metadata(stats: dict[str, Any], action_dim: int) -> dict[str, Any]: + if set(stats) != EXPECTED_NORMALIZATION_PROFILES: + raise StarVLAError( + "unexpected official FAST normalization profiles: " + f"{sorted(stats)}" + ) + metadata: dict[str, Any] = { + "starvla.normalization.profile_count": len(stats), + "starvla.normalization.profile_keys": sorted(stats), + "starvla.normalization.clip_actions": False, + "starvla.normalization.binary_threshold": 0.5, + "starvla.normalization.binary_comparison": "gt", + } + expected_mask = [True] * (action_dim - 1) + [False] + for index, key in enumerate(sorted(stats)): + profile = stats[key] + if not isinstance(profile, dict): + raise StarVLAError(f"normalization profile {key!r} must be an object") + action = profile.get("action") + if not isinstance(action, dict): + raise StarVLAError(f"normalization profile {key!r} has no action object") + for field in ("q01", "q99", "mask"): + values = action.get(field) + if not isinstance(values, list) or len(values) != action_dim: + raise StarVLAError( + f"normalization profile {key!r} action.{field} must " + f"have {action_dim} values" + ) + metadata[f"starvla.normalization.profile.{index}.action_{field}"] = values + q01 = action["q01"] + q99 = action["q99"] + mask = action["mask"] + if any(type(value) is not bool for value in mask) or mask != expected_mask: + raise StarVLAError( + f"normalization profile {key!r} action.mask must be {expected_mask}" + ) + if any( + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not math.isfinite(value) + for value in [*q01, *q99] + ): + raise StarVLAError( + f"normalization profile {key!r} action quantiles must be finite" + ) + if any(q99[axis] < q01[axis] for axis in range(action_dim - 1)): + raise StarVLAError( + f"normalization profile {key!r} has q99 below q01" + ) + metadata[f"starvla.normalization.profile.{index}.key"] = key + + state = profile.get("state") + if not isinstance(state, dict): + raise StarVLAError( + f"normalization profile {key!r} has no state statistics" + ) + state_q01 = state.get("q01") + state_q99 = state.get("q99") + if ( + not isinstance(state_q01, list) + or not isinstance(state_q99, list) + or not state_q01 + or len(state_q01) != len(state_q99) + or any( + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not math.isfinite(value) + for value in [*state_q01, *state_q99] + ) + or any(upper < lower for lower, upper in zip(state_q01, state_q99)) + ): + raise StarVLAError( + f"normalization profile {key!r} has invalid state q01/q99" + ) + metadata[f"starvla.normalization.profile.{index}.state_dimension"] = len( + state_q01 + ) + metadata[f"starvla.normalization.profile.{index}.state_q01"] = state_q01 + metadata[f"starvla.normalization.profile.{index}.state_q99"] = state_q99 + return metadata + + +def _normalize_gguf_metadata_value(value: Any) -> Any: + if isinstance(value, bool) or isinstance(value, str) or value is None: + return value + if isinstance(value, int): + if value < -(2**31) or value >= 2**31: + raise StarVLAError(f"GGUF int32 metadata value is out of range: {value}") + return value + if isinstance(value, float): + if not math.isfinite(value): + raise StarVLAError("GGUF metadata floats must be finite") + return float(np.float32(value)) + if isinstance(value, list): + if not value: + raise StarVLAError("runtime GGUF metadata arrays must not be empty") + return [_normalize_gguf_metadata_value(item) for item in value] + if isinstance(value, dict): + return { + str(key): _normalize_gguf_metadata_value(item) + for key, item in value.items() + } + raise StarVLAError(f"unsupported GGUF metadata value: {value!r}") + + +def build_fast_runtime_policy( + *, + manifest: Mapping[str, Any], + entry: Mapping[str, Any], + codec_entry: Mapping[str, Any], + source_dir: Path, + qwen_dir: Path, + codec_dir: Path, +) -> tuple[dict[str, Any], dict[str, np.ndarray]]: + del entry + source = manifest.get("source") + bundle_uuid = manifest.get("bundle_uuid") + if not isinstance(source, Mapping) or not isinstance(bundle_uuid, str): + raise StarVLAError("FAST staging manifest lacks source/bundle provenance") + + qwen = validate_qwen_config(qwen_dir) + processor = validate_qwen_processor(qwen_dir) + generation = validate_generation_config(qwen_dir) + codec = validate_fast_codec(codec_dir, codec_entry) + effective = effective_fast_config(source_dir) + stats = load_json_object(source_dir / "dataset_statistics.json") + arrays = compile_fast_runtime_tensors(qwen_dir, codec_dir) + offsets = arrays[CODEC_TOKEN_OFFSETS_TENSOR] + token_bytes = arrays[CODEC_TOKEN_BYTES_TENSOR] + if ( + effective.get("cot_prompt") != COT_PROMPT + or effective.get("image_count") != 1 + or effective.get("action_dim") != ACTION_DIM + or effective.get("action_horizon") != ACTION_HORIZON + ): + raise StarVLAError("effective FAST source contract is incompatible") + + metadata: dict[str, Any] = { + "general.architecture": "starvla-policy", + "general.name": "StarVLA Qwen2.5-VL FAST policy", + "general.source.uuid": bundle_uuid, + "starvla.schema_version": 1, + "starvla.framework": FRAMEWORK, + "starvla.model_type": MODEL_TYPE, + "starvla.backbone.arch": BACKBONE, + "starvla.bundle.uuid": bundle_uuid, + "starvla.component.text.filename": TEXT_FILENAME, + "starvla.component.mmproj.filename": MMPROJ_FILENAME, + "starvla.qwen.hidden_size": qwen["text_hidden_size"], + "starvla.qwen.input_embedding_size": qwen["text_hidden_size"], + "starvla.qwen.layer_count": qwen["text_layers"], + "starvla.qwen.vocab_size": qwen["vocab_size"], + "starvla.prompt.cot_template": effective["cot_prompt"], + "starvla.action.dimension": ACTION_DIM, + "starvla.action.horizon": ACTION_HORIZON, + "starvla.action.continuous_dimensions": list(range(ACTION_DIM - 1)), + "starvla.action.binary_dimensions": [ACTION_DIM - 1], + "starvla.image.count": effective["image_count"], + "starvla.image.names": ["image_0"], + "starvla.image.processor_min_pixels": processor["min_pixels"], + "starvla.image.processor_max_pixels": processor["max_pixels"], + "starvla.image.patch_size": processor["patch_size"], + "starvla.image.spatial_merge_size": processor["merge_size"], + "starvla.image.min_token_count": processor["min_image_tokens"], + "starvla.image.max_token_count": processor["max_image_tokens"], + "starvla.fast.generation.max_length": generation["max_length"], + "starvla.fast.generation.eos_token_ids": generation["eos_token_id"], + "starvla.fast.generation.top_k": generation["top_k"], + "starvla.fast.generation.repetition_penalty": generation["repetition_penalty"], + "starvla.fast.action_token.count": ACTION_TOKEN_COUNT, + "starvla.fast.codec.scale": codec["scale"], + "starvla.fast.codec.min_token": codec["min_token"], + "starvla.fast.codec.vocab_size": ACTION_TOKEN_COUNT, + "starvla.fast.codec.time_horizon": ACTION_HORIZON, + "starvla.fast.codec.action_dimension": ACTION_DIM, + "starvla.fast.codec.token_offsets_count": int(offsets.size), + "starvla.fast.codec.token_bytes_count": int(token_bytes.size), + } + metadata.update(normalization_metadata(stats, ACTION_DIM)) + return { + key: _normalize_gguf_metadata_value(value) + for key, value in metadata.items() + }, arrays + + +def _add_runtime_metadata(writer: Any, metadata: Mapping[str, Any]) -> None: + for key in sorted(metadata): + if key == "general.architecture": + continue + value = metadata[key] + if isinstance(value, str): + if not value: + raise StarVLAError(f"GGUF string metadata must be non-empty: {key}") + writer.add_string(key, value) + elif isinstance(value, bool): + writer.add_bool(key, value) + elif isinstance(value, int): + writer.add_int32(key, value) + elif isinstance(value, float): + writer.add_float32(key, value) + elif isinstance(value, list): + if not value: + raise StarVLAError(f"GGUF array metadata must be non-empty: {key}") + writer.add_array(key, value) + else: + raise StarVLAError(f"unsupported GGUF metadata value for {key}: {value!r}") + + +def write_fast_runtime_policy_gguf( + path: Path, + metadata: Mapping[str, Any], + arrays: Mapping[str, np.ndarray], +) -> None: + if path.exists(): + raise StarVLAError(f"refusing to overwrite runtime policy GGUF: {path}") + if set(arrays) != FAST_RUNTIME_TENSOR_NAMES: + raise StarVLAError("FAST runtime GGUF tensor set is incomplete") + path.parent.mkdir(parents=True, exist_ok=True) + writer = gguf.GGUFWriter( + path, + arch="starvla-policy", + use_temp_file=True, + ) + try: + _add_runtime_metadata(writer, metadata) + for name in ( + ACTION_TOKEN_MAP_TENSOR, + CODEC_TOKEN_OFFSETS_TENSOR, + CODEC_TOKEN_BYTES_TENSOR, + ): + array = np.ascontiguousarray(arrays[name]) + writer.add_tensor(name, array) + writer.write_header_to_file() + writer.write_kv_data_to_file() + writer.write_tensors_to_file() + writer.close() + except BaseException: + path.unlink(missing_ok=True) + raise + + +def _gguf_field(reader: Any, key: str) -> Any: + field = reader.get_field(key) + if field is None: + raise StarVLAError(f"FAST runtime GGUF is missing metadata: {key}") + return field.contents() + + +def _metadata_matches(actual: Any, expected: Any) -> bool: + if isinstance(expected, bool): + return isinstance(actual, bool) and actual is expected + if isinstance(expected, float): + return isinstance(actual, (int, float)) and math.isclose( + actual, + expected, + rel_tol=1e-7, + abs_tol=1e-7, + ) + if isinstance(expected, list): + return ( + isinstance(actual, list) + and len(actual) == len(expected) + and all( + _metadata_matches(actual_item, expected_item) + for actual_item, expected_item in zip(actual, expected) + ) + ) + if isinstance(expected, dict): + return ( + isinstance(actual, dict) + and set(actual) == set(expected) + and all( + _metadata_matches(actual[key], expected[key]) + for key in expected + ) + ) + return actual == expected + + +def validate_fast_runtime_policy_gguf( + path: Path, + *, + expected_metadata: Mapping[str, Any] | None = None, + expected_arrays: Mapping[str, np.ndarray] | None = None, +) -> dict[str, Any]: + if not path.is_file() or path.stat().st_size == 0: + raise StarVLAError(f"missing FAST policy GGUF: {path}") + try: + reader = gguf.GGUFReader(path) + except Exception as exc: + raise StarVLAError(f"failed to read FAST policy GGUF: {exc}") from exc + if _gguf_field(reader, "general.architecture") != "starvla-policy": + raise StarVLAError("FAST policy GGUF has the wrong architecture") + + vocab_size = int(_gguf_field(reader, "starvla.fast.codec.vocab_size")) + token_bytes_count = int( + _gguf_field(reader, "starvla.fast.codec.token_bytes_count") + ) + action_dim = int(_gguf_field(reader, "starvla.action.dimension")) + tensors = {tensor.name: tensor for tensor in reader.tensors} + expected_shapes = { + ACTION_TOKEN_MAP_TENSOR: ("I32", [vocab_size]), + CODEC_TOKEN_OFFSETS_TENSOR: ("I32", [vocab_size + 1]), + CODEC_TOKEN_BYTES_TENSOR: ("I8", [token_bytes_count]), + } + if set(tensors) != set(expected_shapes) or len(tensors) != len(reader.tensors): + raise StarVLAError(f"FAST policy tensor set mismatch: {sorted(tensors)}") + for name, (dtype, shape) in expected_shapes.items(): + tensor = tensors[name] + if tensor.tensor_type.name != dtype or list(map(int, tensor.shape)) != shape: + raise StarVLAError(f"FAST policy tensor shape/type mismatch: {name}") + + action_map = np.asarray(tensors[ACTION_TOKEN_MAP_TENSOR].data, dtype=np.int32).reshape(-1) + offsets = np.asarray(tensors[CODEC_TOKEN_OFFSETS_TENSOR].data, dtype=np.int32).reshape(-1) + token_bytes = np.asarray(tensors[CODEC_TOKEN_BYTES_TENSOR].data, dtype=np.int8).reshape(-1) + if ( + np.any(action_map < 0) + or np.unique(action_map).size != vocab_size + or offsets[0] != 0 + or offsets[-1] != token_bytes.size + or np.any(np.diff(offsets) <= 0) + ): + raise StarVLAError("FAST policy codec tensors are invalid") + + profile_count = int(_gguf_field(reader, "starvla.normalization.profile_count")) + profile_keys = _gguf_field(reader, "starvla.normalization.profile_keys") + if profile_count <= 0 or not isinstance(profile_keys, list) or len(profile_keys) != profile_count: + raise StarVLAError("FAST policy normalization profiles are invalid") + for index in range(profile_count): + for suffix in ("action_q01", "action_q99", "action_mask"): + values = _gguf_field(reader, f"starvla.normalization.profile.{index}.{suffix}") + if not isinstance(values, list) or len(values) != action_dim: + raise StarVLAError(f"FAST normalization profile {index} is incomplete") + + if expected_metadata is not None: + expected_keys = set(expected_metadata) + actual_keys = { + key for key in reader.fields + if key.startswith("starvla.") or key in expected_keys + } + if actual_keys != expected_keys: + raise StarVLAError("FAST policy GGUF metadata set mismatch") + for key, expected in expected_metadata.items(): + if not _metadata_matches(_gguf_field(reader, key), expected): + raise StarVLAError(f"FAST policy GGUF metadata mismatch: {key}") + if expected_arrays is not None: + actual_arrays = { + ACTION_TOKEN_MAP_TENSOR: action_map, + CODEC_TOKEN_OFFSETS_TENSOR: offsets, + CODEC_TOKEN_BYTES_TENSOR: token_bytes, + } + if set(expected_arrays) != set(actual_arrays) or any( + not np.array_equal(actual_arrays[name], expected) + for name, expected in expected_arrays.items() + ): + raise StarVLAError("FAST policy tensors differ from compiled assets") + + tensor_contract = { + name: {"dtype": dtype, "shape": shape} + for name, (dtype, shape) in expected_shapes.items() + } + record = { + "path": path.name, + "size": path.stat().st_size, + "sha256": sha256_file(path), + "dtype": "integer_runtime_constants", + "architecture": "starvla-policy", + "tensor_count": len(tensors), + "tensor_dtypes": {"I32": 2, "I8": 1}, + "tensor_contract": tensor_contract, + } + del reader + return record + + +def build_bundle_manifest( + *, + manifest: Mapping[str, Any], + codec: Mapping[str, Any], + text_component: Mapping[str, Any], + mmproj_component: Mapping[str, Any], + policy_component: Mapping[str, Any], +) -> dict[str, Any]: + if policy_component.get("path") != POLICY_FILENAME: + raise StarVLAError("FAST policy component has an unexpected filename") + return { + "schema_version": 1, + "kind": "starvla_qwen25_fast_official_gguf_bundle", + "variant": VARIANT_KEY, + "framework": FRAMEWORK, + "backbone": BACKBONE, + "model_type": MODEL_TYPE, + "bundle_uuid": manifest["bundle_uuid"], + "source": manifest["source"], + "generation": dict(GENERATION_CONTRACT), + "action_token_mapping": manifest["action_token_mapping"], + "fast_codec": { + **dict(codec), + "runtime_storage": "embedded_integer_tensors_in_policy_gguf", + "runtime_policy_gguf": POLICY_FILENAME, + "external_sidecars_required": False, + }, + "components": { + "text": dict(text_component), + "mmproj": dict(mmproj_component), + "policy": dict(policy_component), + }, + "policy_implementation": "finetuned_autoregressive_qwen2_5_vl", + "separate_policy_gguf": POLICY_FILENAME, + } + + +def validate_checkpoint_inventory(records: Sequence[Any]) -> dict[str, Any]: + summary = inventory_summary(list(records)) + mismatches = [ + f"{key}: expected {value!r}, got {summary.get(key)!r}" + for key, value in EXPECTED_INVENTORY.items() + if summary.get(key) != value + ] + by_name = {record.destination_name: record for record in records} + required_shapes = { + "model.embed_tokens.weight": [153713, 2048], + "lm_head.weight": [153713, 2048], + "visual.patch_embed.proj.weight": [1280, 3, 2, 14, 14], + } + for name, shape in required_shapes.items(): + record = by_name.get(name) + if record is None: + mismatches.append(f"{name}: missing") + elif record.shape != shape: + mismatches.append(f"{name}: expected {shape}, got {record.shape}") + if mismatches: + raise StarVLAError( + "Qwen2.5 FAST checkpoint inventory mismatch: " + "; ".join(mismatches) + ) + return summary + + +def preflight( + catalog: Mapping[str, Any], + source_dir: Path, + qwen_dir: Path, + codec_dir: Path, +) -> dict[str, Any]: + entry, qwen_entry, codec_entry = validate_catalog_contract(catalog) + policy_hashes = verify_catalog_files(source_dir, entry) + qwen_hashes = verify_catalog_files(qwen_dir, qwen_entry) + qwen_config = validate_qwen_config(qwen_dir) + qwen_processor = validate_qwen_processor(qwen_dir) + generation = validate_generation_config(qwen_dir) + mapping = validate_action_token_mapping(qwen_dir) + codec = validate_fast_codec(codec_dir, codec_entry) + runtime_arrays = compile_fast_runtime_tensors(qwen_dir, codec_dir) + codec["runtime_tensors"] = { + name: { + "dtype": "I32" if array.dtype == np.int32 else "I8", + "shape": list(array.shape), + "sha256": _sha256_bytes(array.tobytes(order="C")), + } + for name, array in runtime_arrays.items() + } + return { + "variant": VARIANT_KEY, + "framework": FRAMEWORK, + "backbone": BACKBONE, + "model_type": MODEL_TYPE, + "source": { + "repo_id": entry["repo_id"], + "revision": entry["revision"], + "metadata": policy_hashes, + }, + "qwen": { + "repo_id": qwen_entry["repo_id"], + "revision": qwen_entry["revision"], + "metadata": qwen_hashes, + "config": qwen_config, + "processor": qwen_processor, + }, + "generation": generation, + "action_token_mapping": mapping, + "fast_codec": codec, + } + + +def effective_fast_config(source_dir: Path) -> dict[str, Any]: + try: + import yaml + except ImportError as exc: + raise StarVLAError("PyYAML is required to resolve the FAST config") from exc + try: + source = yaml.safe_load((source_dir / "config.yaml").read_text(encoding="utf-8")) + except (OSError, yaml.YAMLError) as exc: + raise StarVLAError(f"failed to load FAST config.yaml: {exc}") from exc + if not isinstance(source, dict): + raise StarVLAError("FAST config.yaml must contain an object") + return { + "schema_version": 1, + "framework": "QwenFast", + "backbone": BACKBONE, + "action_model": "autoregressive_vlm_lm_head", + "action_dim": ACTION_DIM, + "action_horizon": ACTION_HORIZON, + "cot_prompt": COT_PROMPT, + "image_count": 1, + "image_size": [224, 224], + "generation": dict(GENERATION_CONTRACT), + "source_config_sha256": sha256_file(source_dir / "config.yaml"), + "resolved_overrides": { + "framework.action_model.action_model_type": { + "source": source.get("framework", {}) + .get("action_model", {}) + .get("action_model_type"), + "effective": "FAST", + "authority": "pinned_QwenFast_factory", + }, + "framework.action_model.action_horizon": { + "source": None, + "effective": ACTION_HORIZON, + "authority": "future_action_window_size_plus_current_step", + }, + }, + } + + +def stage_checkpoint( + *, + checkpoint: Path, + source_dir: Path, + qwen_dir: Path, + codec_dir: Path, + staging_dir: Path, + catalog: Mapping[str, Any], + max_shard_size: int, + verify_hash: bool, +) -> dict[str, Any]: + entry, qwen_entry, codec_entry = validate_catalog_contract(catalog) + report = preflight(catalog, source_dir, qwen_dir, codec_dir) + if verify_hash: + verify_checkpoint_file(checkpoint, entry) + if staging_dir.exists(): + raise StarVLAError(f"refusing to overwrite staging directory: {staging_dir}") + staging_dir.parent.mkdir(parents=True, exist_ok=True) + staging_dir.mkdir() + try: + state_dict = load_checkpoint_state(checkpoint) + records = build_inventory(state_dict, entry, enforce_expected=False) + inventory = validate_checkpoint_inventory(records) + validate_qwen_vlm_destination_names( + qwen_dir, + qwen_entry, + records, + backbone=BACKBONE, + ) + + hf_dir = staging_dir / "hf" + source_assets_dir = staging_dir / "source" + hf_dir.mkdir() + source_assets_dir.mkdir() + qwen_assets = copy_qwen_assets(qwen_dir, hf_dir, qwen_entry) + expected_qwen_assets = staged_qwen_asset_hashes(qwen_entry) + if qwen_assets != expected_qwen_assets: + raise StarVLAError("staged Qwen asset hashes do not match the catalog overrides") + policy_assets = copy_policy_assets(source_dir, source_assets_dir, entry) + + effective_path = source_assets_dir / "effective_config.json" + atomic_write_json(effective_path, effective_fast_config(source_dir)) + vlm_output = write_safetensor_shards( + hf_dir, + "model", + "model.safetensors.index.json", + list(records), + state_dict, + max_shard_size, + ) + del state_dict + + codec = validate_fast_codec(codec_dir, codec_entry) + manifest = { + "schema_version": 1, + "kind": "starvla_qwen25_fast_official_checkpoint_staging", + "variant": VARIANT_KEY, + "framework": FRAMEWORK, + "backbone": BACKBONE, + "model_type": MODEL_TYPE, + "bundle_uuid": official_bundle_uuid(entry, catalog), + "source": { + "repo_id": entry["repo_id"], + "revision": entry["revision"], + "checkpoint": str(checkpoint.resolve()), + "checkpoint_size": checkpoint.stat().st_size, + "checkpoint_sha256": entry["checkpoint"]["sha256"] + if verify_hash + else sha256_file(checkpoint), + "starvla_revision": catalog["source_revisions"]["starvla"], + "llama_cpp_revision": catalog["source_revisions"]["llama_cpp"], + "qwen_repo_id": qwen_entry["repo_id"], + "qwen_revision": qwen_entry["revision"], + "qwen_asset": QWEN_ASSET_KEY, + }, + "inventory": inventory, + "qwen_assets": qwen_assets, + "policy_assets": policy_assets, + "fast_codec": codec, + "generation": dict(GENERATION_CONTRACT), + "action_token_mapping": report["action_token_mapping"], + "effective_config": { + "path": "source/effective_config.json", + "size": effective_path.stat().st_size, + "sha256": sha256_file(effective_path), + }, + "vlm_output": vlm_output, + "tensors": [record.to_json() for record in records], + } + atomic_write_json( + staging_dir / STAGING_MANIFEST_FILENAME, + manifest, + overwrite=False, + ) + return manifest + except BaseException: + shutil.rmtree(staging_dir, ignore_errors=True) + raise + + +def validate_staging_manifest( + manifest: Mapping[str, Any], + catalog: Mapping[str, Any], + staging_dir: Path, +) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]: + entry, qwen_entry, codec_entry = validate_catalog_contract(catalog) + expected = { + "schema_version": 1, + "kind": "starvla_qwen25_fast_official_checkpoint_staging", + "variant": VARIANT_KEY, + "framework": FRAMEWORK, + "backbone": BACKBONE, + "model_type": MODEL_TYPE, + "bundle_uuid": official_bundle_uuid(entry, catalog), + } + mismatches = [ + f"{key}: expected {value!r}, got {manifest.get(key)!r}" + for key, value in expected.items() + if manifest.get(key) != value + ] + source = manifest.get("source") + if not isinstance(source, Mapping): + mismatches.append("source: missing") + else: + source_expected = { + "repo_id": entry["repo_id"], + "revision": entry["revision"], + "checkpoint_size": entry["checkpoint"]["size"], + "checkpoint_sha256": entry["checkpoint"]["sha256"], + "starvla_revision": catalog["source_revisions"]["starvla"], + "llama_cpp_revision": catalog["source_revisions"]["llama_cpp"], + "qwen_repo_id": qwen_entry["repo_id"], + "qwen_revision": qwen_entry["revision"], + "qwen_asset": QWEN_ASSET_KEY, + } + mismatches.extend( + f"source.{key}: expected {value!r}, got {source.get(key)!r}" + for key, value in source_expected.items() + if source.get(key) != value + ) + inventory = manifest.get("inventory") + if not isinstance(inventory, Mapping): + mismatches.append("inventory: missing") + else: + mismatches.extend( + f"inventory.{key}: expected {value!r}, got {inventory.get(key)!r}" + for key, value in EXPECTED_INVENTORY.items() + if inventory.get(key) != value + ) + if manifest.get("qwen_assets") != staged_qwen_asset_hashes(qwen_entry): + mismatches.append("qwen_assets: mismatch") + expected_policy_assets = { + relative: record["sha256"] for relative, record in entry["file_hashes"].items() + } + if manifest.get("policy_assets") != expected_policy_assets: + mismatches.append("policy_assets: mismatch") + if manifest.get("generation") != GENERATION_CONTRACT: + mismatches.append("generation: mismatch") + expected_mapping = { + "count": ACTION_TOKEN_COUNT, + "fast_token_min": 0, + "fast_token_max": ACTION_TOKEN_COUNT - 1, + "vlm_token_min": ACTION_TOKEN_MIN, + "vlm_token_max": ACTION_TOKEN_MAX, + "mapping": "vlm_token_id = fast_token_id + 151665", + "sha256": qwen_entry["file_hashes"]["added_token_id_map.json"]["sha256"], + } + if manifest.get("action_token_mapping") != expected_mapping: + mismatches.append("action_token_mapping: mismatch") + expected_codec = { + "repo_id": codec_entry["repo_id"], + "revision": codec_entry["revision"], + "scale": 10, + "min_token": -354, + "vocab_size": ACTION_TOKEN_COUNT, + "action_dim": ACTION_DIM, + "time_horizon": ACTION_HORIZON, + "files": { + relative: record["sha256"] + for relative, record in codec_entry["file_hashes"].items() + }, + } + if manifest.get("fast_codec") != expected_codec: + mismatches.append("fast_codec: mismatch") + tensors = manifest.get("tensors") + if not isinstance(tensors, list) or len(tensors) != EXPECTED_INVENTORY["total_tensors"]: + mismatches.append("tensors: incomplete checkpoint inventory") + if mismatches: + raise StarVLAError("invalid Qwen2.5 FAST staging manifest: " + "; ".join(mismatches)) + + hf_dir = staging_dir / "hf" + verify_staged_assets(hf_dir, manifest["qwen_assets"], component="Qwen") + verify_staged_assets( + staging_dir / "source", + manifest["policy_assets"], + component="FAST source", + ) + effective = manifest.get("effective_config") + if not isinstance(effective, Mapping): + raise StarVLAError("FAST staging manifest has no effective_config record") + if effective.get("path") != "source/effective_config.json": + raise StarVLAError("FAST staging manifest has an invalid effective_config path") + effective_path = staging_dir / "source" / "effective_config.json" + if ( + not effective_path.is_file() + or effective_path.stat().st_size != effective.get("size") + or sha256_file(effective_path) != effective.get("sha256") + ): + raise StarVLAError("FAST staged effective_config size/SHA256 mismatch") + index = verify_staged_shards(hf_dir, manifest["vlm_output"], component="FAST VLM") + staged_names = set(index["weight_map"]) + manifest_names = { + str(record.get("destination_name")) + for record in tensors + if isinstance(record, Mapping) + } + if len(manifest_names) != len(tensors) or staged_names != manifest_names: + raise StarVLAError("FAST staged tensor names do not match the checkpoint inventory") + return entry, qwen_entry, codec_entry + + +def _reserve_output_directory(output_dir: Path) -> None: + if output_dir.exists(): + raise StarVLAError(f"refusing to overwrite output directory: {output_dir}") + output_dir.parent.mkdir(parents=True, exist_ok=True) + output_dir.mkdir() + + +def convert_staging( + *, + staging_dir: Path, + source_dir: Path, + qwen_dir: Path, + codec_dir: Path, + output_dir: Path, + catalog: Mapping[str, Any], + llama_root: Path, + python: str, +) -> dict[str, Any]: + manifest_path = staging_dir / STAGING_MANIFEST_FILENAME + try: + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise StarVLAError(f"failed to load staging manifest {manifest_path}: {exc}") from exc + entry, _, codec_entry = validate_staging_manifest(manifest, catalog, staging_dir) + preflight(catalog, source_dir, qwen_dir, codec_dir) + verified_llama = verify_llama_checkout( + llama_root.resolve(strict=True), + str(manifest["source"]["llama_cpp_revision"]), + ) + + _reserve_output_directory(output_dir) + try: + text_output = output_dir / TEXT_FILENAME + mmproj_output = output_dir / MMPROJ_FILENAME + text_metadata = output_dir / "text-metadata.json" + mmproj_metadata = output_dir / "mmproj-metadata.json" + common_metadata = { + "general.source.uuid": manifest["bundle_uuid"], + "general.source.url": ( + f"https://huggingface.co/{entry['repo_id']}/tree/{entry['revision']}" + ), + "general.finetune": "starvla-qwen25-fast", + } + atomic_write_json( + text_metadata, + { + **common_metadata, + "general.name": "StarVLA Qwen2.5-VL FAST text policy", + }, + ) + atomic_write_json( + mmproj_metadata, + { + **common_metadata, + "general.name": "StarVLA Qwen2.5-VL FAST mmproj", + }, + ) + commands = build_commands( + python, + staging_dir / "hf", + text_output, + mmproj_output, + text_metadata, + mmproj_metadata, + "bf16", + "bf16", + llama_root=verified_llama, + ) + for command in commands: + subprocess.run(command, check=True, cwd=REPOSITORY_ROOT) + for output in (text_output, mmproj_output): + if not output.is_file() or output.stat().st_size == 0: + raise StarVLAError(f"converter did not create {output}") + text_metadata.unlink() + mmproj_metadata.unlink() + + policy_metadata, policy_arrays = build_fast_runtime_policy( + manifest=manifest, + entry=entry, + codec_entry=codec_entry, + source_dir=source_dir, + qwen_dir=qwen_dir, + codec_dir=codec_dir, + ) + policy_output = output_dir / POLICY_FILENAME + write_fast_runtime_policy_gguf( + policy_output, + policy_metadata, + policy_arrays, + ) + policy_component = validate_fast_runtime_policy_gguf( + policy_output, + expected_metadata=policy_metadata, + expected_arrays=policy_arrays, + ) + bundle = build_bundle_manifest( + manifest=manifest, + codec=validate_fast_codec(codec_dir, codec_entry), + text_component={ + "path": TEXT_FILENAME, + "size": text_output.stat().st_size, + "sha256": sha256_file(text_output), + "dtype": "bf16", + }, + mmproj_component={ + "path": MMPROJ_FILENAME, + "size": mmproj_output.stat().st_size, + "sha256": sha256_file(mmproj_output), + "dtype": "bf16", + }, + policy_component=policy_component, + ) + atomic_write_json( + output_dir / BUNDLE_MANIFEST_FILENAME, + bundle, + overwrite=False, + ) + return bundle + except BaseException: + shutil.rmtree(output_dir, ignore_errors=True) + raise + + +def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--checkpoint", type=Path) + parser.add_argument("--source-dir", type=Path, required=True) + parser.add_argument("--qwen-assets", type=Path, required=True) + parser.add_argument("--fast-codec", type=Path, required=True) + parser.add_argument("--staging-dir", type=Path, required=True) + parser.add_argument("--output-dir", type=Path) + parser.add_argument("--catalog", type=Path, default=DEFAULT_CATALOG) + parser.add_argument( + "--llama-root", + type=Path, + default=REPOSITORY_ROOT / "third_party" / "llama.cpp", + ) + parser.add_argument("--python", default=sys.executable) + parser.add_argument("--max-shard-size", type=parse_size, default=parse_size("2G")) + parser.add_argument("--preflight", action="store_true") + parser.add_argument("--dry-run", action="store_true") + parser.add_argument("--stage-only", action="store_true") + parser.add_argument("--skip-hash-check", action="store_true") + return parser.parse_args(argv) + + +def main(argv: Sequence[str] | None = None) -> int: + args = parse_args(argv) + try: + catalog = load_catalog(args.catalog) + report = preflight( + catalog, + args.source_dir, + args.qwen_assets, + args.fast_codec, + ) + if args.preflight: + print(json.dumps(report, indent=2, sort_keys=True)) + return 0 + if args.dry_run: + output_dir = args.output_dir or Path("ckpts/starvla/gguf/qwen25-fast") + commands = build_commands( + args.python, + args.staging_dir / "hf", + output_dir / TEXT_FILENAME, + output_dir / MMPROJ_FILENAME, + output_dir / "text-metadata.json", + output_dir / "mmproj-metadata.json", + "bf16", + "bf16", + llama_root=args.llama_root, + ) + print( + json.dumps( + { + **report, + "checkpoint_required_for_execution": True, + "commands": commands, + "bundle_components": { + "text": TEXT_FILENAME, + "mmproj": MMPROJ_FILENAME, + "policy": POLICY_FILENAME, + }, + "runtime_policy": { + "built_in_process": True, + "external_sidecars_required": False, + "tensor_count": len(FAST_RUNTIME_TENSOR_NAMES), + }, + }, + indent=2, + sort_keys=True, + ) + ) + return 0 + if args.checkpoint is None: + raise StarVLAError("--checkpoint is required unless --preflight or --dry-run is used") + manifest = stage_checkpoint( + checkpoint=args.checkpoint, + source_dir=args.source_dir, + qwen_dir=args.qwen_assets, + codec_dir=args.fast_codec, + staging_dir=args.staging_dir, + catalog=catalog, + max_shard_size=args.max_shard_size, + verify_hash=not args.skip_hash_check, + ) + print(f"staging manifest: {args.staging_dir / STAGING_MANIFEST_FILENAME}") + if args.stage_only: + print(json.dumps(manifest["inventory"], indent=2, sort_keys=True)) + return 0 + if args.output_dir is None: + raise StarVLAError("--output-dir is required unless --stage-only is used") + bundle = convert_staging( + staging_dir=args.staging_dir, + source_dir=args.source_dir, + qwen_dir=args.qwen_assets, + codec_dir=args.fast_codec, + output_dir=args.output_dir, + catalog=catalog, + llama_root=args.llama_root, + python=args.python, + ) + print(f"bundle manifest: {args.output_dir / BUNDLE_MANIFEST_FILENAME}") + print(json.dumps(bundle["components"], indent=2, sort_keys=True)) + return 0 + except ( + StarVLAError, + OSError, + json.JSONDecodeError, + subprocess.CalledProcessError, + KeyError, + ) as exc: + print(f"error: {exc}", file=sys.stderr) + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/hf2gguf/starvla/generate_starvla_qwen25_fast_golden.py b/tools/hf2gguf/starvla/generate_starvla_qwen25_fast_golden.py new file mode 100644 index 0000000..4bdd889 --- /dev/null +++ b/tools/hf2gguf/starvla/generate_starvla_qwen25_fast_golden.py @@ -0,0 +1,765 @@ +#!/usr/bin/env python3 +"""Generate a CUDA local-Python action golden from the official Qwen2.5 FAST .pt.""" + +from __future__ import annotations + +import argparse +import hashlib +import importlib.metadata +import json +import math +import os +import platform +import random +import sys +from pathlib import Path +from typing import Any, Mapping, Sequence + + +TOOLS_DIR = Path(__file__).resolve().parent +if str(TOOLS_DIR) not in sys.path: + sys.path.insert(0, str(TOOLS_DIR)) + +from convert_starvla_qwen25_fast import ( # noqa: E402 + ACTION_DIM, + ACTION_HORIZON, + ACTION_TOKEN_MAX, + ACTION_TOKEN_MIN, + BACKBONE, + COT_PROMPT, + FAST_CODEC_ASSET_KEY, + FRAMEWORK, + GENERATION_CONTRACT, + MODEL_TYPE, + STAGING_MANIFEST_FILENAME, + VARIANT_KEY, + validate_fast_codec, + validate_staging_manifest, +) +from starvla_checkpoint import ( # noqa: E402 + DEFAULT_CATALOG, + StarVLAError, + atomic_write_json, + get_variant, + load_catalog, + official_bundle_uuid, + sha256_file, + verify_checkpoint_file, +) + + +SCHEMA_VERSION = 2 +GOLDEN_KIND = "starvla_qwen25_fast_local_python_action_golden" +DEFAULT_SEED = 42 +UNNORM_KEYS = ("bridge_dataset", "fractal20220817_data") +EXPECTED_RUNTIME_VERSIONS = { + "torch": "2.6.0", + "torchvision": "0.21.0", + "transformers": "4.57.0", + "numpy": "1.26.4", + "qwen-vl-utils": "0.0.14", +} + + +def canonical_sha256(value: Any) -> str: + payload = json.dumps(value, sort_keys=True, separators=(",", ":")).encode("utf-8") + return hashlib.sha256(payload).hexdigest() + + +def valid_sha256(value: Any) -> bool: + return ( + isinstance(value, str) + and len(value) == 64 + and all(character in "0123456789abcdef" for character in value) + ) + + +def validate_runtime_versions(actual: Mapping[str, str]) -> None: + mismatches = [] + for name, expected in EXPECTED_RUNTIME_VERSIONS.items(): + version = str(actual.get(name, "missing")).split("+", 1)[0] + if version != expected: + mismatches.append(f"{name}: expected {expected}, got {version}") + if mismatches: + raise StarVLAError( + "Qwen2.5 FAST local-Python runtime version mismatch: " + + "; ".join(mismatches) + ) + + +def distribution_version(name: str) -> str: + try: + return importlib.metadata.version(name) + except importlib.metadata.PackageNotFoundError: + return "missing" + + +def render_prompt(instruction: str) -> str: + if not isinstance(instruction, str) or not instruction.strip(): + raise StarVLAError("FAST instruction must be a non-empty string") + return COT_PROMPT.replace("{instruction}", instruction) + + +def build_messages(image: Any, instruction: str) -> list[dict[str, Any]]: + return [ + { + "role": "user", + "content": [ + {"type": "image", "image": image}, + {"type": "text", "text": render_prompt(instruction)}, + ], + } + ] + + +def extract_action_token_ids( + generated_ids: Sequence[Sequence[int]], +) -> list[list[int]]: + result = [] + for row in generated_ids: + tokens = [] + for value in row: + if not isinstance(value, int) or isinstance(value, bool): + raise StarVLAError("generated token IDs must be integers") + if ACTION_TOKEN_MIN <= value <= ACTION_TOKEN_MAX: + tokens.append(value) + result.append(tokens) + return result + + +def map_vlm_to_fast_ids( + batch_action_token_ids: Sequence[Sequence[int]], +) -> list[list[int]]: + result = [] + for row in batch_action_token_ids: + fast_ids = [token_id - ACTION_TOKEN_MIN for token_id in row] + if any(token_id < 0 or token_id > 2047 for token_id in fast_ids): + raise StarVLAError("generated action token is outside the FAST vocabulary") + result.append(fast_ids) + return result + + +def validate_actions(value: Any, *, name: str) -> list[list[list[float]]]: + try: + import numpy as np + except ImportError as exc: + raise StarVLAError("NumPy is required for FAST action validation") from exc + actions = np.asarray(value, dtype=np.float64) + expected_shape = (1, ACTION_HORIZON, ACTION_DIM) + if actions.shape != expected_shape: + raise StarVLAError( + f"{name} has shape {list(actions.shape)}, expected {list(expected_shape)}" + ) + if not np.isfinite(actions).all(): + raise StarVLAError(f"{name} contains a non-finite value") + return actions.tolist() + + +def validate_normalized_actions(value: Any) -> list[list[list[float]]]: + return validate_actions(value, name="FAST normalized actions") + + +def load_normalization_profile( + dataset_statistics: Path, + unnorm_key: str, +) -> dict[str, Any]: + if unnorm_key not in UNNORM_KEYS: + raise StarVLAError( + f"FAST --unnorm-key must be one of {list(UNNORM_KEYS)}, got {unnorm_key!r}" + ) + try: + statistics = json.loads(dataset_statistics.read_text(encoding="utf-8")) + except (OSError, UnicodeError, json.JSONDecodeError) as exc: + raise StarVLAError( + f"failed to load FAST dataset statistics {dataset_statistics}: {exc}" + ) from exc + if not isinstance(statistics, dict) or set(statistics) != set(UNNORM_KEYS): + raise StarVLAError("FAST dataset statistics profile set is incompatible") + try: + action = statistics[unnorm_key]["action"] + q01 = [float(value) for value in action["q01"]] + q99 = [float(value) for value in action["q99"]] + mask = list(action["mask"]) + except (KeyError, TypeError, ValueError) as exc: + raise StarVLAError("FAST action normalization statistics are malformed") from exc + if ( + len(q01) != ACTION_DIM + or len(q99) != ACTION_DIM + or mask != [True] * 6 + [False] + or not all(math.isfinite(value) for value in q01 + q99) + or any(high < low for low, high in zip(q01[:6], q99[:6])) + ): + raise StarVLAError("FAST action normalization profile is incompatible") + return { + "profile": unnorm_key, + "action_q01": q01, + "action_q99": q99, + "action_mask": mask, + "continuous_dimensions": [0, 1, 2, 3, 4, 5], + "binary_dimensions": [6], + "binary_threshold": 0.5, + "binary_comparison": "gt", + "clip_actions": False, + } + + +def unnormalize_actions( + normalized_actions: Any, + normalization: Mapping[str, Any], +) -> list[list[list[float]]]: + try: + import numpy as np + except ImportError as exc: + raise StarVLAError("NumPy is required for FAST action unnormalization") from exc + normalized = np.asarray( + validate_actions(normalized_actions, name="FAST normalized actions"), + dtype=np.float32, + ) + q01 = np.asarray(normalization["action_q01"], dtype=np.float32) + q99 = np.asarray(normalization["action_q99"], dtype=np.float32) + result = np.empty_like(normalized, dtype=np.float32) + result[..., :6] = (normalized[..., :6] + np.float32(1.0)) * np.float32( + 0.5 + ) * (q99[:6] - q01[:6]) + q01[:6] + result[..., 6] = (normalized[..., 6] > np.float32(0.5)).astype(np.float32) + return validate_actions(result.tolist(), name="FAST unnormalized actions") + + +def validate_fast_token_rows( + fast_processor: Any, + batch_fast_ids: Sequence[Sequence[int]], +) -> None: + expected_coefficients = ACTION_HORIZON * ACTION_DIM + for index, row in enumerate(batch_fast_ids): + try: + decoded = fast_processor.bpe_tokenizer.decode(list(row)) + except Exception as exc: + raise StarVLAError(f"FAST token row {index} cannot be decoded") from exc + if len(decoded) != expected_coefficients: + raise StarVLAError( + f"FAST token row {index} contains {len(decoded)} coefficients; " + f"expected {expected_coefficients}" + ) + + +def finalize_golden_id(value: dict[str, Any]) -> str: + payload = dict(value) + payload.pop("golden_id", None) + golden_id = canonical_sha256(payload) + value["golden_id"] = golden_id + return golden_id + + +def _require_regular_bound_file( + path_value: Any, + size_value: Any, + sha_value: Any, + *, + label: str, +) -> Path: + path = Path(str(path_value)) + if not path.is_absolute() or not path.is_file() or path.is_symlink(): + raise StarVLAError(f"{label} must be an absolute regular file") + if ( + not isinstance(size_value, int) + or isinstance(size_value, bool) + or path.stat().st_size != size_value + or not valid_sha256(sha_value) + or sha256_file(path) != sha_value + ): + raise StarVLAError(f"{label} no longer matches its bound size/SHA256") + return path + + +def validate_golden( + value: Any, + *, + verify_files: bool = False, + catalog_path: Path = DEFAULT_CATALOG, +) -> dict[str, Any]: + if not isinstance(value, dict) or value.get("kind") != GOLDEN_KIND: + raise StarVLAError("not a Qwen2.5 FAST local-Python golden") + catalog = load_catalog(catalog_path) + variant = get_variant(catalog, VARIANT_KEY) + qwen = catalog["shared_assets"][variant["qwen_asset"]] + codec = catalog["shared_assets"][FAST_CODEC_ASSET_KEY] + expected = { + "schema_version": SCHEMA_VERSION, + "variant": VARIANT_KEY, + "framework": FRAMEWORK, + "backbone": BACKBONE, + "model_type": MODEL_TYPE, + "bundle_uuid": official_bundle_uuid(variant, catalog), + "generation": GENERATION_CONTRACT, + } + mismatches = [ + f"{key}: expected {item!r}, got {value.get(key)!r}" + for key, item in expected.items() + if value.get(key) != item + ] + source = value.get("source") + input_record = value.get("input") + prompt = value.get("prompt") + normalization = value.get("normalization") + runtime = value.get("runtime") + result = value.get("result") + if not all( + isinstance(item, dict) + for item in ( + source, + input_record, + prompt, + normalization, + runtime, + result, + ) + ): + mismatches.append( + "source/input/prompt/normalization/runtime/result must be objects" + ) + if mismatches: + raise StarVLAError("invalid Qwen2.5 FAST golden: " + "; ".join(mismatches)) + + source_expected = { + "checkpoint_repo_id": variant["repo_id"], + "checkpoint_revision": variant["revision"], + "checkpoint_filename": Path(variant["checkpoint"]["path"]).name, + "checkpoint_size": variant["checkpoint"]["size"], + "checkpoint_sha256": variant["checkpoint"]["sha256"], + "qwen_repo_id": qwen["repo_id"], + "qwen_revision": qwen["revision"], + "fast_codec_repo_id": codec["repo_id"], + "fast_codec_revision": codec["revision"], + "fast_codec_files": { + relative: codec["file_hashes"][relative]["sha256"] + for relative in codec["files"] + }, + "starvla_revision": catalog["source_revisions"]["starvla"], + "llama_revision": catalog["source_revisions"]["llama_cpp"], + "weight_source": "official_original_pt_staged_exact_weights", + } + for key, expected_value in source_expected.items(): + if source.get(key) != expected_value: + mismatches.append(f"source.{key}: expected {expected_value!r}") + if source.get("bundle_uuid") != value["bundle_uuid"]: + mismatches.append("source.bundle_uuid") + if prompt != { + "chat_template_sha256": qwen["file_hashes"]["chat_template.jinja"][ + "sha256" + ] + }: + mismatches.append("prompt.chat_template_sha256") + + instruction = input_record.get("instruction") + if ( + not isinstance(instruction, str) + or input_record.get("framework_prompt") != render_prompt(instruction) + or input_record.get("unnorm_key") not in UNNORM_KEYS + or not valid_sha256(input_record.get("image_sha256")) + or input_record.get("prompt_length", 0) <= 0 + or canonical_sha256(input_record.get("input_ids")) + != input_record.get("input_ids_sha256") + ): + mismatches.append("input prompt/image/profile/token contract") + + expected_profile = load_normalization_profile( + Path(source["dataset_statistics_path"]), + input_record["unnorm_key"], + ) + for key, expected_value in expected_profile.items(): + if normalization.get(key) != expected_value: + mismatches.append(f"normalization.{key}") + if ( + source.get("dataset_statistics_sha256") + != variant["file_hashes"]["dataset_statistics.json"]["sha256"] + or normalization.get("source_sha256") + != source.get("dataset_statistics_sha256") + ): + mismatches.append("normalization source SHA256") + + validate_runtime_versions(runtime) + if ( + runtime.get("backend") != "cuda" + or runtime.get("full_gpu_model") is not True + or runtime.get("dtype") != "bfloat16" + or runtime.get("attn_implementation") != "sdpa" + or runtime.get("tf32") is not False + ): + mismatches.append("runtime CUDA/dtype/attention contract") + + generated = result.get("generated_ids") + action_ids = result.get("action_token_ids") + fast_ids = result.get("fast_token_ids") + if not isinstance(generated, list) or len(generated) != 1: + mismatches.append("result.generated_ids") + elif extract_action_token_ids(generated) != action_ids: + mismatches.append("result.action_token_ids") + if isinstance(action_ids, list): + try: + if map_vlm_to_fast_ids(action_ids) != fast_ids: + mismatches.append("result.fast_token_ids") + except StarVLAError as exc: + mismatches.append(str(exc)) + if ( + not isinstance(action_ids, list) + or len(action_ids) != 1 + or not action_ids[0] + or not isinstance(fast_ids, list) + or len(fast_ids) != 1 + or not fast_ids[0] + ): + mismatches.append("result requires non-empty action/FAST token IDs") + for key in ("normalized_actions", "unnormalized_actions"): + try: + validate_actions(result.get(key), name=f"result.{key}") + except StarVLAError as exc: + mismatches.append(str(exc)) + if result.get("unnormalized_actions") != unnormalize_actions( + result.get("normalized_actions"), normalization + ): + mismatches.append("result.unnormalized_actions formula") + if ( + canonical_sha256(result.get("generated_ids")) + != result.get("generated_ids_sha256") + or canonical_sha256(result.get("normalized_actions")) + != result.get("normalized_actions_sha256") + or canonical_sha256(result.get("unnormalized_actions")) + != result.get("unnormalized_actions_sha256") + ): + mismatches.append("result canonical SHA256") + golden_id = value.get("golden_id") + payload = dict(value) + payload.pop("golden_id", None) + if not valid_sha256(golden_id) or canonical_sha256(payload) != golden_id: + mismatches.append("golden_id") + if mismatches: + raise StarVLAError("invalid Qwen2.5 FAST golden: " + "; ".join(mismatches)) + + if verify_files: + checkpoint = _require_regular_bound_file( + source["checkpoint_path"], + source["checkpoint_size"], + source["checkpoint_sha256"], + label="golden source checkpoint", + ) + verify_checkpoint_file(checkpoint, variant) + statistics_path = _require_regular_bound_file( + source["dataset_statistics_path"], + source["dataset_statistics_size"], + source["dataset_statistics_sha256"], + label="golden dataset statistics", + ) + if statistics_path.name != "dataset_statistics.json": + raise StarVLAError("golden dataset statistics filename is incompatible") + image_path = _require_regular_bound_file( + input_record["image_path"], + input_record["image_size"], + input_record["image_sha256"], + label="golden input image", + ) + if image_path.resolve() != Path(input_record["image_path"]): + raise StarVLAError("golden image path is not canonical") + manifest_path = _require_regular_bound_file( + source["staging_manifest_path"], + source["staging_manifest_size"], + source["staging_manifest_sha256"], + label="golden staging manifest", + ) + try: + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise StarVLAError(f"failed to parse golden staging manifest: {exc}") from exc + staging_dir = Path(source["staged_hf_path"]).parent + validate_staging_manifest(manifest, catalog, staging_dir) + if ( + Path(source["staged_hf_path"]) != (staging_dir / "hf").resolve() + or Path(manifest["source"]["checkpoint"]).resolve() != checkpoint.resolve() + or manifest["bundle_uuid"] != value["bundle_uuid"] + ): + raise StarVLAError("golden staged exact-weight binding is inconsistent") + codec_path = Path(str(source["fast_codec_path"])) + if ( + not codec_path.is_absolute() + or not codec_path.is_dir() + or codec_path.is_symlink() + ): + raise StarVLAError("golden FAST codec path is not a bound directory") + actual_codec = validate_fast_codec(codec_path, codec) + if actual_codec["files"] != source["fast_codec_files"]: + raise StarVLAError("golden FAST codec source binding changed") + return value + + +def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--staged-hf", type=Path, required=True) + parser.add_argument("--staging-manifest", type=Path) + parser.add_argument("--fast-codec", type=Path, required=True) + parser.add_argument("--image", type=Path, required=True) + parser.add_argument("--instruction", required=True) + parser.add_argument("--unnorm-key", choices=UNNORM_KEYS, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--catalog", type=Path, default=DEFAULT_CATALOG) + parser.add_argument( + "--attn-implementation", + choices=("sdpa",), + default="sdpa", + ) + parser.add_argument("--seed", type=int, default=DEFAULT_SEED) + return parser.parse_args(argv) + + +def main(argv: Sequence[str] | None = None) -> int: + args = parse_args(argv) + try: + if not sys.flags.isolated: + raise StarVLAError( + "Qwen2.5 FAST golden generation must run in isolated mode (`python -I`)" + ) + os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" + import numpy as np + import torch + import transformers + from PIL import Image + from qwen_vl_utils import process_vision_info + from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration + + if not torch.cuda.is_available(): + raise StarVLAError("Qwen2.5 FAST golden generation requires CUDA") + runtime_versions = { + "torch": torch.__version__, + "torchvision": distribution_version("torchvision"), + "transformers": transformers.__version__, + "numpy": np.__version__, + "qwen-vl-utils": distribution_version("qwen-vl-utils"), + } + validate_runtime_versions(runtime_versions) + catalog = load_catalog(args.catalog) + variant = get_variant(catalog, VARIANT_KEY) + qwen = catalog["shared_assets"][variant["qwen_asset"]] + codec_entry = catalog["shared_assets"][FAST_CODEC_ASSET_KEY] + verify_checkpoint_file(args.checkpoint, variant) + codec = validate_fast_codec(args.fast_codec.resolve(), codec_entry) + + manifest_path = ( + args.staging_manifest + or args.staged_hf.parent / STAGING_MANIFEST_FILENAME + ).resolve() + try: + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise StarVLAError( + f"failed to load staging manifest {manifest_path}: {exc}" + ) from exc + staging_dir = args.staged_hf.resolve().parent + validate_staging_manifest(manifest, catalog, staging_dir) + if ( + args.staged_hf.resolve() != staging_dir / "hf" + or Path(manifest["source"]["checkpoint"]).resolve() + != args.checkpoint.resolve() + ): + raise StarVLAError( + "staging manifest is bound to a different exact-weight source" + ) + dataset_statistics = args.checkpoint.resolve().parents[1] / "dataset_statistics.json" + normalization = load_normalization_profile( + dataset_statistics, args.unnorm_key + ) + normalization["source_sha256"] = sha256_file(dataset_statistics) + expected_stats_sha = variant["file_hashes"]["dataset_statistics.json"]["sha256"] + if normalization["source_sha256"] != expected_stats_sha: + raise StarVLAError("official FAST dataset statistics SHA256 changed") + if not args.image.is_file() or args.image.is_symlink(): + raise StarVLAError(f"missing FAST golden image: {args.image}") + if args.seed < 0: + raise StarVLAError("--seed must be non-negative") + + random.seed(args.seed) + np.random.seed(args.seed) + torch.manual_seed(args.seed) + torch.cuda.manual_seed_all(args.seed) + torch.use_deterministic_algorithms(True) + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + torch.backends.cudnn.benchmark = False + + processor = AutoProcessor.from_pretrained( + args.staged_hf, local_files_only=True + ) + processor.tokenizer.padding_side = "left" + model = Qwen2_5_VLForConditionalGeneration.from_pretrained( + args.staged_hf, + local_files_only=True, + torch_dtype=torch.bfloat16, + attn_implementation=args.attn_implementation, + ).to("cuda") + model.eval() + + generation_file = json.loads( + (args.staged_hf / "generation_config.json").read_text(encoding="utf-8") + ) + for key, expected in GENERATION_CONTRACT.items(): + if key != "max_length" and generation_file.get(key) != expected: + raise StarVLAError( + f"staged generation_config drift at {key}: " + f"expected {expected!r}, got {generation_file.get(key)!r}" + ) + + with Image.open(args.image) as opened: + image = opened.convert("RGB") + messages = build_messages(image, args.instruction) + rendered_chat = processor.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True + ) + image_inputs, video_inputs = process_vision_info([messages]) + inputs = processor( + text=[rendered_chat], + images=image_inputs, + videos=video_inputs, + padding=True, + return_tensors="pt", + ).to("cuda") + input_ids = inputs["input_ids"].detach().cpu().tolist() + prompt_length = int(inputs["input_ids"].shape[1]) + + with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16): + generated_tensor = model.generate( + **inputs, max_length=GENERATION_CONTRACT["max_length"] + ) + generated_ids = generated_tensor.detach().cpu().tolist() + action_token_ids = extract_action_token_ids(generated_ids) + fast_token_ids = map_vlm_to_fast_ids(action_token_ids) + + fast_processor = AutoProcessor.from_pretrained( + args.fast_codec, + trust_remote_code=True, + local_files_only=True, + ) + fast_processor.time_horizon = ACTION_HORIZON + fast_processor.action_dim = ACTION_DIM + decode_inputs = [row if row else None for row in fast_token_ids] + validate_fast_token_rows(fast_processor, fast_token_ids) + normalized_actions = validate_normalized_actions( + fast_processor.decode(decode_inputs) + ) + unnormalized = unnormalize_actions(normalized_actions, normalization) + + checkpoint = args.checkpoint.resolve() + image_path = args.image.resolve() + golden: dict[str, Any] = { + "schema_version": SCHEMA_VERSION, + "kind": GOLDEN_KIND, + "variant": VARIANT_KEY, + "framework": FRAMEWORK, + "backbone": BACKBONE, + "model_type": MODEL_TYPE, + "bundle_uuid": manifest["bundle_uuid"], + "source": { + "bundle_uuid": manifest["bundle_uuid"], + "checkpoint_repo_id": variant["repo_id"], + "checkpoint_revision": variant["revision"], + "checkpoint_filename": checkpoint.name, + "checkpoint_path": str(checkpoint), + "checkpoint_size": checkpoint.stat().st_size, + "checkpoint_sha256": variant["checkpoint"]["sha256"], + "dataset_statistics_path": str(dataset_statistics.resolve()), + "dataset_statistics_size": dataset_statistics.stat().st_size, + "dataset_statistics_sha256": expected_stats_sha, + "qwen_repo_id": qwen["repo_id"], + "qwen_revision": qwen["revision"], + "fast_codec_repo_id": codec["repo_id"], + "fast_codec_revision": codec["revision"], + "fast_codec_files": codec["files"], + "starvla_revision": catalog["source_revisions"]["starvla"], + "llama_revision": catalog["source_revisions"]["llama_cpp"], + "staged_hf_path": str(args.staged_hf.resolve()), + "staging_manifest_path": str(manifest_path), + "staging_manifest_size": manifest_path.stat().st_size, + "staging_manifest_sha256": sha256_file(manifest_path), + "fast_codec_path": str(args.fast_codec.resolve()), + "weight_source": "official_original_pt_staged_exact_weights", + }, + "prompt": { + "chat_template_sha256": qwen["file_hashes"][ + "chat_template.jinja" + ]["sha256"], + }, + "runtime": { + "python": platform.python_version(), + **runtime_versions, + "backend": "cuda", + "full_gpu_model": True, + "device": torch.cuda.get_device_name(torch.cuda.current_device()), + "cuda": torch.version.cuda, + "dtype": "bfloat16", + "attn_implementation": args.attn_implementation, + "tf32": False, + "torch_deterministic_algorithms": True, + "cublas_workspace_config": ":4096:8", + "seed": args.seed, + }, + "input": { + "image_path": str(image_path), + "image_size": image_path.stat().st_size, + "image_sha256": sha256_file(image_path), + "decoded_size": [image.width, image.height], + "instruction": args.instruction, + "unnorm_key": args.unnorm_key, + "framework_prompt": render_prompt(args.instruction), + "rendered_chat_template": rendered_chat, + "input_ids": input_ids, + "input_ids_sha256": canonical_sha256(input_ids), + "prompt_length": prompt_length, + }, + "normalization": normalization, + "generation": dict(GENERATION_CONTRACT), + "result": { + "generated_ids": generated_ids, + "continuation_ids": [ + row[prompt_length:] for row in generated_ids + ], + "generated_ids_sha256": canonical_sha256(generated_ids), + "action_token_ids": action_token_ids, + "fast_token_ids": fast_token_ids, + "normalized_actions": normalized_actions, + "normalized_actions_sha256": canonical_sha256(normalized_actions), + "unnormalized_actions": unnormalized, + "unnormalized_actions_sha256": canonical_sha256(unnormalized), + }, + } + finalize_golden_id(golden) + validate_golden( + golden, verify_files=True, catalog_path=args.catalog + ) + atomic_write_json(args.output, golden, overwrite=False) + print(f"Qwen2.5 FAST golden: {args.output}") + print( + json.dumps( + { + "golden_id": golden["golden_id"], + "prompt_tokens": prompt_length, + "generated_tokens": len(generated_ids[0]) - prompt_length, + "action_tokens": len(action_token_ids[0]), + "action_shape": [1, ACTION_HORIZON, ACTION_DIM], + "unnorm_key": args.unnorm_key, + }, + indent=2, + sort_keys=True, + ) + ) + return 0 + except ( + StarVLAError, + OSError, + ValueError, + RuntimeError, + KeyError, + json.JSONDecodeError, + ) as exc: + print(f"error: {exc}", file=sys.stderr) + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) From 39f29f7b45cd132ce871600a5370bc49b3f2aac8 Mon Sep 17 00:00:00 2001 From: JJJYmmm <1650675829@qq.com> Date: Mon, 10 Aug 2026 12:42:57 +0800 Subject: [PATCH 09/11] starvla: integrate seven Qwen policy variants --- CMakeLists.txt | 127 +- patches/llama.cpp/README.md | 23 +- robot_client/cpp/model_client.cpp | 5 +- robot_client/cpp/model_client.h | 1 + robot_client/python/model_client.py | 14 +- robot_server/model-server.cpp | 93 +- robot_server/model_adapter.cpp | 5 +- robot_server/protocol.cpp | 50 +- robot_server/protocol.h | 3 +- .../shell/launch_robot_server_linux_cuda.sh | 18 +- robot_server/test/benchmark_latency.py | 16 +- robot_server/test/test_server_latency.sh | 25 +- src/model-cli.cpp | 90 +- src/models/argument_parse.h | 27 + src/models/ggml_backend.cpp | 2 +- src/models/ggml_backend.h | 2 + src/models/model.h | 9 + src/models/model_factory.cpp | 13 +- src/models/model_type.cpp | 49 + src/models/pi0/pi0_model.cpp | 4 + src/models/smolvla/smolvla_model.cpp | 4 + src/models/starvla/fast_codec.cpp | 475 +---- src/models/starvla/fast_codec.h | 63 +- src/models/starvla/fast_policy.cpp | 369 +--- src/models/starvla/fast_policy.h | 35 +- src/models/starvla/groot_policy.cpp | 566 ++--- src/models/starvla/groot_policy.h | 53 +- src/models/starvla/groot_prompt.cpp | 35 +- src/models/starvla/groot_prompt.h | 13 +- src/models/starvla/normalization.cpp | 33 +- src/models/starvla/normalization.h | 3 +- src/models/starvla/oft_image_preprocess.cpp | 198 +- src/models/starvla/oft_image_preprocess.h | 30 +- src/models/starvla/oft_policy.cpp | 325 +-- src/models/starvla/oft_policy.h | 29 +- src/models/starvla/oft_prompt.cpp | 28 +- src/models/starvla/oft_prompt.h | 17 +- src/models/starvla/pi_policy.cpp | 1002 +++------ src/models/starvla/pi_policy.h | 66 +- src/models/starvla/pi_v3_policy.cpp | 663 ++---- src/models/starvla/pi_v3_policy.h | 66 +- src/models/starvla/policy_gguf.h | 144 ++ src/models/starvla/qwen3vl_bridge.cpp | 859 ++++---- src/models/starvla/qwen3vl_bridge.h | 100 +- src/models/starvla/qwen_bf16_round_cuda.cu | 218 ++ src/models/starvla/qwen_bf16_round_cuda.h | 28 + src/models/starvla/starvla_engine.cpp | 987 +++++++++ src/models/starvla/starvla_engine.h | 74 + src/models/starvla/starvla_model.cpp | 93 + src/models/starvla/starvla_model.h | 35 + tests/starvla/fast_codec_test.cpp | 194 -- tests/starvla/fast_runtime_test.cpp | 137 -- tests/starvla/groot_prompt_test.cpp | 86 - tests/starvla/model_test.cpp | 59 + tests/starvla/oft_image_preprocess_test.cpp | 162 -- tests/starvla/oft_prompt_test.cpp | 81 - tests/starvla/test_starvla_pi_v3_golden.py | 36 - tests/starvla/test_starvla_qwen25_fast.py | 111 - ...ly_starvla_patches.sh => apply_patches.sh} | 10 +- tools/hf2gguf/README.md | 3 +- tools/hf2gguf/README_ZH.md | 3 +- tools/hf2gguf/starvla/README.md | 177 ++ tools/hf2gguf/starvla/checkpoint_catalog.json | 478 +++++ .../starvla/compare_starvla_actions.py | 82 - tools/hf2gguf/starvla/convert.sh | 160 ++ tools/hf2gguf/starvla/convert_starvla_all.sh | 144 ++ .../starvla/convert_starvla_policy_to_gguf.py | 1120 ++++------ .../starvla/convert_starvla_qwen25_fast.py | 200 +- .../starvla/convert_starvla_qwen_to_gguf.py | 21 +- tools/hf2gguf/starvla/download_starvla.py | 126 +- tools/hf2gguf/starvla/environment.yaml | 15 +- .../starvla/generate_starvla_groot_golden.py | 1169 ----------- .../starvla/generate_starvla_oft_golden.py | 928 --------- .../starvla/generate_starvla_pi_v3_golden.py | 1831 ----------------- .../generate_starvla_qwen25_fast_golden.py | 765 ------- .../generate_starvla_qwen25_groot_golden.py | 1142 ---------- .../generate_starvla_qwen25_oft_golden.py | 986 --------- .../generate_starvla_qwen25_pi_golden.py | 1338 ------------ .../starvla/groot_golden_constraints.txt | 13 - .../starvla/inspect_starvla_checkpoint.py | 177 -- .../starvla/pi_v3_golden_constraints.txt | 14 - .../starvla/serve_starvla_groot_reference.py | 535 ----- .../starvla/serve_starvla_oft_reference.py | 1152 ----------- tools/hf2gguf/starvla/starvla_checkpoint.py | 287 ++- tools/hf2gguf/starvla/starvla_surgery.py | 32 +- .../hf2gguf/starvla/starvla_variant_config.sh | 55 + .../starvla/validate_starvla_bundle.py | 271 +-- 87 files changed, 5453 insertions(+), 15834 deletions(-) create mode 100644 src/models/argument_parse.h create mode 100644 src/models/model_type.cpp create mode 100644 src/models/starvla/policy_gguf.h create mode 100644 src/models/starvla/qwen_bf16_round_cuda.cu create mode 100644 src/models/starvla/qwen_bf16_round_cuda.h create mode 100644 src/models/starvla/starvla_engine.cpp create mode 100644 src/models/starvla/starvla_engine.h create mode 100644 src/models/starvla/starvla_model.cpp create mode 100644 src/models/starvla/starvla_model.h delete mode 100644 tests/starvla/fast_codec_test.cpp delete mode 100644 tests/starvla/fast_runtime_test.cpp delete mode 100644 tests/starvla/groot_prompt_test.cpp create mode 100644 tests/starvla/model_test.cpp delete mode 100644 tests/starvla/oft_image_preprocess_test.cpp delete mode 100644 tests/starvla/oft_prompt_test.cpp delete mode 100644 tests/starvla/test_starvla_pi_v3_golden.py delete mode 100644 tests/starvla/test_starvla_qwen25_fast.py rename tools/{llama_cpp/apply_starvla_patches.sh => apply_patches.sh} (91%) create mode 100644 tools/hf2gguf/starvla/README.md create mode 100644 tools/hf2gguf/starvla/checkpoint_catalog.json delete mode 100644 tools/hf2gguf/starvla/compare_starvla_actions.py create mode 100755 tools/hf2gguf/starvla/convert.sh create mode 100755 tools/hf2gguf/starvla/convert_starvla_all.sh delete mode 100644 tools/hf2gguf/starvla/generate_starvla_groot_golden.py delete mode 100644 tools/hf2gguf/starvla/generate_starvla_oft_golden.py delete mode 100644 tools/hf2gguf/starvla/generate_starvla_pi_v3_golden.py delete mode 100644 tools/hf2gguf/starvla/generate_starvla_qwen25_fast_golden.py delete mode 100644 tools/hf2gguf/starvla/generate_starvla_qwen25_groot_golden.py delete mode 100644 tools/hf2gguf/starvla/generate_starvla_qwen25_oft_golden.py delete mode 100644 tools/hf2gguf/starvla/generate_starvla_qwen25_pi_golden.py delete mode 100644 tools/hf2gguf/starvla/groot_golden_constraints.txt delete mode 100755 tools/hf2gguf/starvla/inspect_starvla_checkpoint.py delete mode 100644 tools/hf2gguf/starvla/pi_v3_golden_constraints.txt delete mode 100644 tools/hf2gguf/starvla/serve_starvla_groot_reference.py delete mode 100644 tools/hf2gguf/starvla/serve_starvla_oft_reference.py create mode 100644 tools/hf2gguf/starvla/starvla_variant_config.sh diff --git a/CMakeLists.txt b/CMakeLists.txt index 28174de..7c6c01b 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1,9 +1,12 @@ cmake_minimum_required(VERSION 3.16) project(robotcpp VERSION 0.1.0 LANGUAGES C CXX) +include(CTest) + option(ROBOT_CPP_BUILD_ROBOT_SERVER "Build model-server target" ON) option(ROBOT_CPP_BUILD_MODEL_CLI "Build model-cli target" OFF) option(ROBOT_CPP_BUILD_ROBOT_CLIENT "Build C++ robot client targets" OFF) +option(ROBOT_CPP_BUILD_STARVLA "Build the StarVLA runtime (requires llama.cpp overlay)" OFF) set(CMAKE_CXX_STANDARD 17) set(CMAKE_CXX_STANDARD_REQUIRED ON) @@ -12,6 +15,27 @@ set(CMAKE_CXX_EXTENSIONS OFF) if(NOT EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/third_party/llama.cpp/CMakeLists.txt") message(FATAL_ERROR "third_party/llama.cpp is required; run `git submodule update --init --recursive`") endif() +if(ROBOT_CPP_BUILD_STARVLA) + file(READ + "${CMAKE_CURRENT_SOURCE_DIR}/third_party/llama.cpp/tools/mtmd/models/qwen3vl.cpp" + ROBOT_CPP_QWEN3VL_MTMD_SOURCE) + file(READ + "${CMAKE_CURRENT_SOURCE_DIR}/third_party/llama.cpp/include/llama.h" + ROBOT_CPP_LLAMA_PUBLIC_HEADER) + string(FIND "${ROBOT_CPP_QWEN3VL_MTMD_SOURCE}" + "FFN_GELU_ERF" ROBOT_CPP_QWEN3VL_PARITY_PATCH_INDEX) + string(FIND "${ROBOT_CPP_LLAMA_PUBLIC_HEADER}" + "llama_set_backend_native_graphs_enabled" ROBOT_CPP_LLAMA_GRAPH_PATCH_INDEX) + if(ROBOT_CPP_QWEN3VL_PARITY_PATCH_INDEX EQUAL -1 OR + ROBOT_CPP_LLAMA_GRAPH_PATCH_INDEX EQUAL -1) + message(FATAL_ERROR + "StarVLA requires the pinned llama.cpp overlay. Run " + "`./tools/apply_patches.sh` from the repository root, " + "then configure again.") + endif() + unset(ROBOT_CPP_QWEN3VL_MTMD_SOURCE) + unset(ROBOT_CPP_LLAMA_PUBLIC_HEADER) +endif() set(LLAMA_BUILD_COMMON ON CACHE BOOL "" FORCE) set(LLAMA_BUILD_TOOLS OFF CACHE BOOL "" FORCE) set(LLAMA_BUILD_EXAMPLES OFF CACHE BOOL "" FORCE) @@ -20,7 +44,16 @@ add_subdirectory(third_party/llama.cpp EXCLUDE_FROM_ALL) if(NOT TARGET ggml OR NOT TARGET llama) message(FATAL_ERROR "llama.cpp must provide ggml and llama targets") endif() - +if(ROBOT_CPP_BUILD_STARVLA) + # mtmd normally inherits this variable when llama.cpp builds all tools. + if(NOT LLAMA_INSTALL_VERSION) + set(LLAMA_INSTALL_VERSION ${PROJECT_VERSION}) + endif() + add_subdirectory(third_party/llama.cpp/tools/mtmd EXCLUDE_FROM_ALL) + if(NOT TARGET mtmd) + message(FATAL_ERROR "llama.cpp must provide the mtmd target for Qwen-VL") + endif() +endif() set(ROBOT_CPP_LLAMA_INCLUDE_DIRS ${CMAKE_CURRENT_SOURCE_DIR}/third_party/llama.cpp ${CMAKE_CURRENT_SOURCE_DIR}/third_party/llama.cpp/include @@ -30,14 +63,29 @@ set(ROBOT_CPP_LLAMA_INCLUDE_DIRS ) set(SMOLVLA_DIR ${CMAKE_CURRENT_SOURCE_DIR}/src/models/smolvla) +set(STARVLA_DIR ${CMAKE_CURRENT_SOURCE_DIR}/src/models/starvla) set(ROBOT_SERVER_DIR ${CMAKE_CURRENT_SOURCE_DIR}/robot_server) set(ROBOT_CLIENT_DIR ${CMAKE_CURRENT_SOURCE_DIR}/robot_client) -add_library(smolvla_runtime STATIC +add_library(robotcpp_model_common STATIC src/models/ggml_backend.cpp src/models/ggml_backend.h src/models/gguf_loader.cpp src/models/gguf_loader.h + src/models/model_type.cpp +) +target_include_directories(robotcpp_model_common + PUBLIC + ${CMAKE_CURRENT_SOURCE_DIR}/src + ${ROBOT_CPP_LLAMA_INCLUDE_DIRS} +) +target_link_libraries(robotcpp_model_common PUBLIC ggml) +target_compile_features(robotcpp_model_common PUBLIC cxx_std_17) +if(NOT MSVC) + target_compile_options(robotcpp_model_common PRIVATE -Wno-cast-qual) +endif() + +add_library(smolvla_runtime STATIC ${SMOLVLA_DIR}/smolvla_engine.cpp ${SMOLVLA_DIR}/smolvla_engine.h ${SMOLVLA_DIR}/state_proj.cpp @@ -53,17 +101,13 @@ target_include_directories(smolvla_runtime ${SMOLVLA_DIR} ${ROBOT_CPP_LLAMA_INCLUDE_DIRS} ) -target_link_libraries(smolvla_runtime PUBLIC ggml llama) +target_link_libraries(smolvla_runtime PUBLIC robotcpp_model_common llama) target_compile_features(smolvla_runtime PUBLIC cxx_std_17) if(NOT MSVC) target_compile_options(smolvla_runtime PRIVATE -Wno-cast-qual) endif() add_library(pi0_engine STATIC - src/models/ggml_backend.cpp - src/models/ggml_backend.h - src/models/gguf_loader.cpp - src/models/gguf_loader.h src/models/pi0/types.h src/models/pi0/action.cpp src/models/pi0/action.h @@ -87,13 +131,65 @@ target_include_directories(pi0_engine ${CMAKE_CURRENT_SOURCE_DIR}/src ${ROBOT_CPP_LLAMA_INCLUDE_DIRS} ) -target_link_libraries(pi0_engine PUBLIC ggml llama) +target_link_libraries(pi0_engine PUBLIC robotcpp_model_common llama) target_compile_features(pi0_engine PUBLIC cxx_std_17) if(NOT MSVC) target_compile_options(pi0_engine PRIVATE -Wno-cast-qual) endif() +if(ROBOT_CPP_BUILD_STARVLA) +add_library(starvla_runtime STATIC + ${STARVLA_DIR}/fast_codec.cpp + ${STARVLA_DIR}/fast_codec.h + ${STARVLA_DIR}/fast_policy.cpp + ${STARVLA_DIR}/fast_policy.h + ${STARVLA_DIR}/groot_policy.cpp + ${STARVLA_DIR}/groot_policy.h + ${STARVLA_DIR}/groot_prompt.cpp + ${STARVLA_DIR}/groot_prompt.h + ${STARVLA_DIR}/normalization.cpp + ${STARVLA_DIR}/normalization.h + ${STARVLA_DIR}/oft_image_preprocess.cpp + ${STARVLA_DIR}/oft_image_preprocess.h + ${STARVLA_DIR}/oft_prompt.cpp + ${STARVLA_DIR}/oft_prompt.h + ${STARVLA_DIR}/oft_policy.cpp + ${STARVLA_DIR}/oft_policy.h + ${STARVLA_DIR}/pi_policy.cpp + ${STARVLA_DIR}/pi_policy.h + ${STARVLA_DIR}/pi_v3_policy.cpp + ${STARVLA_DIR}/pi_v3_policy.h + ${STARVLA_DIR}/policy_gguf.h + ${STARVLA_DIR}/qwen3vl_bridge.cpp + ${STARVLA_DIR}/qwen3vl_bridge.h + ${STARVLA_DIR}/starvla_engine.cpp + ${STARVLA_DIR}/starvla_engine.h + third_party/llama.cpp/examples/gguf-hash/deps/sha256/sha256.c +) +target_include_directories(starvla_runtime + PUBLIC + ${CMAKE_CURRENT_SOURCE_DIR}/src + ${ROBOT_CPP_LLAMA_INCLUDE_DIRS} + PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/third_party/llama.cpp/vendor + ${CMAKE_CURRENT_SOURCE_DIR}/third_party/llama.cpp/examples/gguf-hash/deps + ${CMAKE_CURRENT_SOURCE_DIR}/third_party/llama.cpp/examples/gguf-hash/deps/sha256 +) +target_link_libraries(starvla_runtime PUBLIC robotcpp_model_common llama mtmd) +target_compile_features(starvla_runtime PUBLIC cxx_std_17) +if(GGML_CUDA) + enable_language(CUDA) + target_sources(starvla_runtime PRIVATE ${STARVLA_DIR}/qwen_bf16_round_cuda.cu) + target_compile_definitions(starvla_runtime PRIVATE ROBOTCPP_STARVLA_CUDA=1) + set_property(TARGET starvla_runtime PROPERTY CUDA_STANDARD 17) +endif() +if(NOT MSVC) + target_compile_options(starvla_runtime PRIVATE -Wno-cast-qual) +endif() +endif() + add_library(robotcpp STATIC + src/models/argument_parse.h src/models/model.h src/models/model_factory.cpp src/models/pi0/pi0_model.cpp @@ -108,6 +204,19 @@ target_include_directories(robotcpp ) target_link_libraries(robotcpp PUBLIC smolvla_runtime pi0_engine) target_compile_features(robotcpp PUBLIC cxx_std_17) +if(ROBOT_CPP_BUILD_STARVLA) + target_sources(robotcpp PRIVATE + ${STARVLA_DIR}/starvla_model.cpp + ${STARVLA_DIR}/starvla_model.h) + target_link_libraries(robotcpp PUBLIC starvla_runtime) + target_compile_definitions(robotcpp PUBLIC ROBOT_CPP_BUILD_STARVLA=1) +endif() + +if(BUILD_TESTING AND ROBOT_CPP_BUILD_STARVLA) + add_executable(robotcpp-starvla-model-test tests/starvla/model_test.cpp) + target_link_libraries(robotcpp-starvla-model-test PRIVATE robotcpp) + add_test(NAME robotcpp-starvla-model-test COMMAND robotcpp-starvla-model-test) +endif() if(ROBOT_CPP_BUILD_ROBOT_SERVER OR ROBOT_CPP_BUILD_ROBOT_CLIENT) add_library(robot_server_common STATIC @@ -139,6 +248,7 @@ if(ROBOT_CPP_BUILD_ROBOT_CLIENT) set_target_properties(model-cpp-client-example PROPERTIES RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/bin) target_link_libraries(model-cpp-client-example PRIVATE model_client_cpp) target_compile_features(model-cpp-client-example PRIVATE cxx_std_17) + endif() if(ROBOT_CPP_BUILD_ROBOT_SERVER) @@ -165,7 +275,6 @@ if(ROBOT_CPP_BUILD_ROBOT_SERVER) ) target_link_libraries(model-server PRIVATE robot_server_core robotcpp) target_compile_features(model-server PRIVATE cxx_std_17) - add_executable(smolvla-raw-predict ${ROBOT_SERVER_DIR}/test/smolvla_raw_predict.cpp) set_target_properties(smolvla-raw-predict PROPERTIES RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/bin) target_include_directories(smolvla-raw-predict PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src ${SMOLVLA_DIR}) diff --git a/patches/llama.cpp/README.md b/patches/llama.cpp/README.md index 6601890..1a17011 100644 --- a/patches/llama.cpp/README.md +++ b/patches/llama.cpp/README.md @@ -1,21 +1,20 @@ -# StarVLA llama.cpp patches +# llama.cpp patches The project pins `third_party/llama.cpp` at commit -`3e941b813b1acbbf06c2203a94ceb33d84748c1e`. The StarVLA runtime needs two -small changes that are not available through that revision's public APIs: +`3e941b813b1acbbf06c2203a94ceb33d84748c1e`. The repository applies two +changes that are not available through that revision's public APIs: -1. `0001-qwen3vl-vision-parity.patch` matches the upstream Qwen3-VL reference - implementation's position interpolation and exact GELU operations. These - changes are required for action-value parity with the original checkpoint. +1. `0001-qwen3vl-vision-parity.patch` uses the position interpolation and exact + GELU operations from the Qwen3-VL implementation used by StarVLA. 2. `0002-per-context-native-graph-control.patch` adds an optional backend API to disable CUDA graph capture for the text and vision contexts owned by one - StarVLA instance. It prevents retained CUDA graphs from violating long-loop - memory stability gates without globally changing other llama.cpp users. + StarVLA instance. This avoids retained CUDA graphs growing memory use during + long runs without changing the setting for other llama.cpp users. -Apply both patches before configuring or building the StarVLA runtime: +Apply the repository patch set after initializing submodules and before building: ```bash -./tools/llama_cpp/apply_starvla_patches.sh +./tools/apply_patches.sh ``` The command verifies the exact llama.cpp revision and refuses a dirty or @@ -24,8 +23,8 @@ partially patched checkout. It is safe to run again after a complete apply. Inspect or remove the overlay with: ```bash -./tools/llama_cpp/apply_starvla_patches.sh --check -./tools/llama_cpp/apply_starvla_patches.sh --revert +./tools/apply_patches.sh --check +./tools/apply_patches.sh --revert ``` The parent repository commits only these patch assets. It does not advance or diff --git a/robot_client/cpp/model_client.cpp b/robot_client/cpp/model_client.cpp index 5496144..3ad82d8 100644 --- a/robot_client/cpp/model_client.cpp +++ b/robot_client/cpp/model_client.cpp @@ -55,8 +55,9 @@ bool make_predict_request(const ModelObservation & obs, proto::predict_request & return false; } - req.task = obs.prompt; - req.state = obs.state; + req.task = obs.prompt; + req.state = obs.state; + req.initial_noise = obs.initial_noise; req.images.clear(); req.images.reserve(obs.images.size()); diff --git a/robot_client/cpp/model_client.h b/robot_client/cpp/model_client.h index a6b7730..ea84f47 100644 --- a/robot_client/cpp/model_client.h +++ b/robot_client/cpp/model_client.h @@ -21,6 +21,7 @@ struct ModelImage { struct ModelObservation { std::vector images; std::vector state; + std::vector initial_noise; std::string prompt = "grab the block."; }; diff --git a/robot_client/python/model_client.py b/robot_client/python/model_client.py index 06e3791..5c79423 100644 --- a/robot_client/python/model_client.py +++ b/robot_client/python/model_client.py @@ -9,7 +9,7 @@ MAGIC = 0x414C5653 -VERSION = 3 +VERSION = 4 HEADER_SIZE = 32 OP_HEALTH = 1 @@ -21,8 +21,8 @@ IMAGE_RAW_RGB_U8 = 1 HEADER = struct.Struct(" tuple[int, int, int, bytes]: def encode_predict_observation(observation: dict[str, Any]) -> bytes: images = observation["images"] state = state_to_list(observation["state"]) + initial_noise = state_to_list(observation.get("initial_noise")) prompt = str(observation["prompt"]) if not images: raise ValueError("observation.images must contain at least one image") @@ -88,13 +89,14 @@ def encode_predict_observation(observation: dict[str, Any]) -> bytes: encoded_images.append((name, rgb, width, height, stride)) payload = bytearray() - payload += PREDICT_REQ_V2_FIXED.pack( + payload += PREDICT_REQ_FIXED.pack( len(encoded_images), len(state), + len(initial_noise), len(prompt_bytes), ) for name, rgb, width, height, stride in encoded_images: - payload += PREDICT_REQ_V2_IMAGE.pack( + payload += PREDICT_REQ_IMAGE.pack( IMAGE_RAW_RGB_U8, len(name), width, @@ -105,6 +107,8 @@ def encode_predict_observation(observation: dict[str, Any]) -> bytes: ) for value in state: payload += struct.pack(" @@ -26,6 +27,7 @@ struct server_args { std::string action_decoder_path; std::string state_proj_path; std::string action_expert_path; + std::string policy_path; std::string task = "grab the block."; std::string host = "127.0.0.1"; int port = 5555; @@ -37,18 +39,6 @@ struct server_args { int verbosity = 0; }; -static bool parse_model_type(const std::string & value, robotcpp::model_type & out) { - if (value == "smolvla") { - out = robotcpp::model_type::smolvla; - return true; - } - if (value == "pi0") { - out = robotcpp::model_type::pi0; - return true; - } - return false; -} - static bool parse_noise_mode(const std::string & value, int & out_mode) { if (value == "gaussian") { out_mode = SMOLVLA_NOISE_MODE_GAUSSIAN; @@ -74,9 +64,11 @@ static void print_usage(const char * prog) { " [options]\n" " %s --model-type pi0 --vit --mmproj --llm --tokenizer --state-gguf " " --action-decoder [options]\n" + " %s --model-type starvla --llm --mmproj --policy [options]\n" "\n" "Common options:\n" - " --model-type Model type (default: smolvla)\n" + " --model-type smolvla|pi0|starvla\n" + " (default: smolvla)\n" "\n" "SmolVLA options:\n" " --llm LLM GGUF path\n" @@ -93,6 +85,11 @@ static void print_usage(const char * prog) { " --state-gguf State projector GGUF path\n" " --action-decoder Action decoder GGUF path\n" "\n" + "StarVLA options:\n" + " --policy StarVLA policy GGUF path (required)\n" + " --llm Qwen text GGUF path (required)\n" + " --mmproj Qwen vision GGUF path (required)\n" + "\n" "Runtime options:\n" " --host Listen host (default: 127.0.0.1)\n" " --port Listen port (default: 5555)\n" @@ -103,7 +100,7 @@ static void print_usage(const char * prog) { " --noise-seed RNG seed, <0 means auto (default: -1)\n" " --verbosity Log verbosity (default: 0)\n" " -h, --help Show this help\n", - prog, prog); + prog, prog, prog); } // TODO: may need to be cleaned up and optimized @@ -116,10 +113,12 @@ static bool parse_args(int argc, char ** argv, server_args & args) { } else if (arg == "--llm" && i + 1 < argc) { args.llm_path = argv[++i]; } else if (arg == "--model-type" && i + 1 < argc) { - if (!parse_model_type(argv[++i], args.model_type)) { + if (!robotcpp::parse_model_type(argv[++i], args.model_type)) { std::fprintf(stderr, "Error: unsupported model type '%s'\n", argv[i]); return false; } + } else if (arg == "--policy" && i + 1 < argc) { + args.policy_path = argv[++i]; } else if (arg == "--mmproj" && i + 1 < argc) { args.mmproj_path = argv[++i]; } else if (arg == "--vit" && i + 1 < argc) { @@ -139,22 +138,46 @@ static bool parse_args(int argc, char ** argv, server_args & args) { } else if (arg == "--host" && i + 1 < argc) { args.host = argv[++i]; } else if (arg == "--port" && i + 1 < argc) { - args.port = std::atoi(argv[++i]); + const char * value = argv[++i]; + if (!robotcpp::parse_integer_argument(value, args.port)) { + std::fprintf(stderr, "Error: invalid --port value '%s'\n", value); + return false; + } } else if (arg == "--threads" && i + 1 < argc) { - args.threads = std::atoi(argv[++i]); + const char * value = argv[++i]; + if (!robotcpp::parse_integer_argument(value, args.threads)) { + std::fprintf(stderr, "Error: invalid --threads value '%s'\n", value); + return false; + } } else if (arg == "--n-batch" && i + 1 < argc) { - args.n_batch = std::atoi(argv[++i]); + const char * value = argv[++i]; + if (!robotcpp::parse_integer_argument(value, args.n_batch)) { + std::fprintf(stderr, "Error: invalid --n-batch value '%s'\n", value); + return false; + } } else if (arg == "--n-ctx" && i + 1 < argc) { - args.n_ctx = std::atoi(argv[++i]); + const char * value = argv[++i]; + if (!robotcpp::parse_integer_argument(value, args.n_ctx)) { + std::fprintf(stderr, "Error: invalid --n-ctx value '%s'\n", value); + return false; + } } else if (arg == "--noise-mode" && i + 1 < argc) { if (!parse_noise_mode(argv[++i], args.noise_mode)) { std::fprintf(stderr, "Error: invalid noise mode '%s'\n", argv[i]); return false; } } else if (arg == "--noise-seed" && i + 1 < argc) { - args.noise_seed = (int64_t)std::atoll(argv[++i]); + const char * value = argv[++i]; + if (!robotcpp::parse_integer_argument(value, args.noise_seed)) { + std::fprintf(stderr, "Error: invalid --noise-seed value '%s'\n", value); + return false; + } } else if (arg == "--verbosity" && i + 1 < argc) { - args.verbosity = std::atoi(argv[++i]); + const char * value = argv[++i]; + if (!robotcpp::parse_integer_argument(value, args.verbosity)) { + std::fprintf(stderr, "Error: invalid --verbosity value '%s'\n", value); + return false; + } } else { std::fprintf(stderr, "Error: unknown argument '%s'\n", arg.c_str()); return false; @@ -168,15 +191,36 @@ static bool parse_args(int argc, char ** argv, server_args & args) { std::fprintf(stderr, "Error: model-server only listens on 127.0.0.1 in this phase\n"); return false; } + if (args.threads < 0 || args.n_batch <= 0 || args.n_ctx <= 0 || args.verbosity < 0) { + std::fprintf(stderr, + "Error: --threads/--verbosity must be non-negative and --n-batch/--n-ctx must be positive\n"); + return false; + } + if (robotcpp::is_starvla_model_type(args.model_type) && args.noise_mode != SMOLVLA_NOISE_MODE_GAUSSIAN) { + std::fprintf(stderr, "Error: StarVLA does not support --noise-mode debug-sin; use Gaussian noise\n"); + return false; + } if (args.model_type == robotcpp::model_type::smolvla) { if (args.llm_path.empty() || args.mmproj_path.empty() || args.state_proj_path.empty() || args.action_expert_path.empty()) { std::fprintf(stderr, "Error: smolvla requires --llm --mmproj --state-proj --action-expert\n"); return false; } - } else if (args.vit_path.empty() || args.mmproj_path.empty() || args.llm_path.empty() || - args.tokenizer_path.empty() || args.state_path.empty() || args.action_decoder_path.empty()) { - std::fprintf(stderr, "Error: pi0 requires --vit --mmproj --llm --tokenizer --state-gguf --action-decoder\n"); + } else if (args.model_type == robotcpp::model_type::pi0) { + if (args.vit_path.empty() || args.mmproj_path.empty() || args.llm_path.empty() || args.tokenizer_path.empty() || + args.state_path.empty() || args.action_decoder_path.empty()) { + std::fprintf(stderr, + "Error: pi0 requires --vit --mmproj --llm --tokenizer --state-gguf --action-decoder\n"); + return false; + } + } else if (robotcpp::is_starvla_model_type(args.model_type)) { + if (args.llm_path.empty() || args.mmproj_path.empty() || args.policy_path.empty()) { + std::fprintf(stderr, "Error: %s requires --llm --mmproj --policy\n", + robotcpp::model_type_name(args.model_type)); + return false; + } + } else { + std::fprintf(stderr, "Error: unsupported model type '%s'\n", robotcpp::model_type_name(args.model_type)); return false; } return true; @@ -195,6 +239,7 @@ static robotcpp::model_args make_model_args(const server_args & args) { model_args.action_decoder_path = args.action_decoder_path; model_args.state_proj_path = args.state_proj_path; model_args.action_expert_path = args.action_expert_path; + model_args.policy_path = args.policy_path; model_args.n_batch = args.n_batch; model_args.n_ctx = args.n_ctx; model_args.noise_mode = args.noise_mode; diff --git a/robot_server/model_adapter.cpp b/robot_server/model_adapter.cpp index 31556ad..4b1a7b8 100644 --- a/robot_server/model_adapter.cpp +++ b/robot_server/model_adapter.cpp @@ -30,8 +30,9 @@ bool model_adapter::predict(const proto::predict_request & req, proto::predict_r image.stride_bytes = static_cast(src.stride_bytes); obs.images.push_back(image); } - obs.state = req.state; - obs.task = req.task; + obs.state = req.state; + obs.initial_noise = req.initial_noise; + obs.task = req.task; robotcpp::model_result result; if (!model_->predict(obs, result, error)) { diff --git a/robot_server/protocol.cpp b/robot_server/protocol.cpp index 8ea4c48..930a7db 100644 --- a/robot_server/protocol.cpp +++ b/robot_server/protocol.cpp @@ -1,5 +1,6 @@ #include "protocol.h" +#include #include #include @@ -122,6 +123,16 @@ static bool checked_u32_count(size_t n, const char * label, std::string & error) return true; } +static bool validate_f32_array(const std::vector & values, const char * label, std::string & error) { + for (float value : values) { + if (!std::isfinite(value)) { + error = std::string(label) + " contains a non-finite value"; + return false; + } + } + return true; +} + static bool validate_image_payload(const image_payload & image, const char * label, std::string & error) { if (image.image_format != image_raw_rgb_u8) { error = std::string(label) + " unsupported image format"; @@ -216,7 +227,13 @@ bool encode_predict_request(const predict_request & req, std::vector & return false; } if (!checked_u32_count(req.images.size(), "images", error) || - !checked_u32_count(req.state.size(), "state", error) || !checked_u32_count(req.task.size(), "task", error)) { + !checked_u32_count(req.state.size(), "state", error) || + !checked_u32_count(req.initial_noise.size(), "initial noise", error) || + !checked_u32_count(req.task.size(), "task", error)) { + return false; + } + if (!validate_f32_array(req.state, "state", error) || + !validate_f32_array(req.initial_noise, "initial noise", error)) { return false; } for (size_t i = 0; i < req.images.size(); ++i) { @@ -228,6 +245,7 @@ bool encode_predict_request(const predict_request & req, std::vector & put_u32(out, (uint32_t)req.images.size()); put_u32(out, (uint32_t)req.state.size()); + put_u32(out, (uint32_t)req.initial_noise.size()); put_u32(out, (uint32_t)req.task.size()); for (const image_payload & image : req.images) { put_u32(out, image.image_format); @@ -242,6 +260,9 @@ bool encode_predict_request(const predict_request & req, std::vector & for (float v : req.state) { put_f32(out, v); } + for (float v : req.initial_noise) { + put_f32(out, v); + } out.insert(out.end(), req.task.begin(), req.task.end()); for (const image_payload & image : req.images) { out.insert(out.end(), image.name.begin(), image.name.end()); @@ -255,9 +276,10 @@ bool decode_predict_request(const std::vector & payload, predict_reques reader r(payload.data(), payload.size()); uint32_t image_count = 0; uint32_t state_dim = 0; + uint32_t noise_dim = 0; uint32_t task_len = 0; - if (!r.u32(image_count) || !r.u32(state_dim) || !r.u32(task_len)) { + if (!r.u32(image_count) || !r.u32(state_dim) || !r.u32(noise_dim) || !r.u32(task_len)) { error = "short predict request"; return false; } @@ -265,6 +287,11 @@ bool decode_predict_request(const std::vector & payload, predict_reques error = "predict request requires at least one image"; return false; } + constexpr size_t image_metadata_size = 6 * sizeof(uint32_t) + sizeof(uint64_t); + if (image_count > r.remaining() / image_metadata_size) { + error = "image count exceeds predict request metadata"; + return false; + } req.images.assign(image_count, image_payload{}); std::vector name_lens(image_count, 0); @@ -286,6 +313,12 @@ bool decode_predict_request(const std::vector & payload, predict_reques } } + const uint64_t scalar_bytes = (static_cast(state_dim) + static_cast(noise_dim)) * sizeof(float); + if (scalar_bytes > r.remaining() || task_len > r.remaining() - scalar_bytes) { + error = "predict request fields exceed payload"; + return false; + } + req.state.assign(state_dim, 0.0f); for (uint32_t i = 0; i < state_dim; ++i) { if (!r.f32(req.state[i])) { @@ -293,6 +326,19 @@ bool decode_predict_request(const std::vector & payload, predict_reques return false; } } + if (!validate_f32_array(req.state, "state", error)) { + return false; + } + req.initial_noise.assign(noise_dim, 0.0f); + for (uint32_t i = 0; i < noise_dim; ++i) { + if (!r.f32(req.initial_noise[i])) { + error = "short initial noise array"; + return false; + } + } + if (!validate_f32_array(req.initial_noise, "initial noise", error)) { + return false; + } if (!r.string(req.task, task_len)) { error = "short task string"; return false; diff --git a/robot_server/protocol.h b/robot_server/protocol.h index 97ae387..a5f483e 100644 --- a/robot_server/protocol.h +++ b/robot_server/protocol.h @@ -9,7 +9,7 @@ namespace robot_server { namespace protocol { static constexpr uint32_t k_magic = 0x414c5653u; // "SVLA" in little-endian bytes. -static constexpr uint16_t k_version = 3; +static constexpr uint16_t k_version = 4; static constexpr uint16_t k_header_size = 32; static constexpr uint64_t k_default_max_payload = 256ull * 1024ull * 1024ull; @@ -63,6 +63,7 @@ struct metric { struct predict_request { std::vector images; std::vector state; + std::vector initial_noise; std::string task; }; diff --git a/robot_server/shell/launch_robot_server_linux_cuda.sh b/robot_server/shell/launch_robot_server_linux_cuda.sh index 990ed44..c6002b9 100755 --- a/robot_server/shell/launch_robot_server_linux_cuda.sh +++ b/robot_server/shell/launch_robot_server_linux_cuda.sh @@ -21,6 +21,10 @@ SKIP_BUILD="${SKIP_BUILD:-0}" CMAKE_BIN="${CMAKE_BIN:-cmake}" GGML_NATIVE="${GGML_NATIVE:-OFF}" GGML_OPENMP="${GGML_OPENMP:-OFF}" +ROBOT_CPP_BUILD_STARVLA="${ROBOT_CPP_BUILD_STARVLA:-OFF}" +if [ "${MODEL_TYPE}" = "starvla" ]; then + ROBOT_CPP_BUILD_STARVLA=ON +fi SERVER_BIN="${BUILD_DIR}/bin/model-server" @@ -32,7 +36,8 @@ if [ "${SKIP_BUILD}" != "1" ]; then -DGGML_OPENMP="${GGML_OPENMP}" \ -DGGML_CUDA=ON \ -DGGML_METAL=OFF \ - -DROBOT_CPP_BUILD_ROBOT_SERVER=ON + -DROBOT_CPP_BUILD_ROBOT_SERVER=ON \ + -DROBOT_CPP_BUILD_STARVLA="${ROBOT_CPP_BUILD_STARVLA}" echo "== build ==" "${CMAKE_BIN}" --build "${BUILD_DIR}" --target model-server -j8 @@ -70,6 +75,17 @@ case "${MODEL_TYPE}" in --action-decoder "${ACTION_DECODER_GGUF}" ) ;; + starvla) + LLM_GGUF="${LLM_GGUF:?LLM_GGUF must be set for StarVLA}" + MMPROJ_GGUF="${MMPROJ_GGUF:?MMPROJ_GGUF must be set for StarVLA}" + POLICY_GGUF="${POLICY_GGUF:?POLICY_GGUF must be set for StarVLA}" + MODEL_ARGS=( + --model-type starvla + --llm "${LLM_GGUF}" + --mmproj "${MMPROJ_GGUF}" + --policy "${POLICY_GGUF}" + ) + ;; *) echo "unsupported MODEL_TYPE=${MODEL_TYPE}" >&2 exit 1 diff --git a/robot_server/test/benchmark_latency.py b/robot_server/test/benchmark_latency.py index e000113..fa2ef72 100644 --- a/robot_server/test/benchmark_latency.py +++ b/robot_server/test/benchmark_latency.py @@ -26,8 +26,15 @@ def make_random_state(dim: int, seed: int) -> np.ndarray: return rng.uniform(-1.0, 1.0, size=(dim,)).astype(np.float32) -def make_random_observation(width: int, height: int, state_dim: int, prompt: str, image_names: list[str]) -> dict: - return { +def make_random_observation( + width: int, + height: int, + state_dim: int, + initial_noise_dim: int, + prompt: str, + image_names: list[str], +) -> dict: + observation = { "images": [ { "name": image_name, @@ -38,6 +45,9 @@ def make_random_observation(width: int, height: int, state_dim: int, prompt: str "state": make_random_state(state_dim, seed=1), "prompt": prompt, } + if initial_noise_dim: + observation["initial_noise"] = make_random_state(initial_noise_dim, seed=2) + return observation def ordered_columns(rows: list[dict[str, float]]) -> list[str]: @@ -141,6 +151,7 @@ def main() -> int: parser.add_argument("--height", type=int, default=224) parser.add_argument("--image-name", action="append") parser.add_argument("--state-dim", type=int, default=6) + parser.add_argument("--initial-noise-dim", type=int, default=0) parser.add_argument("--prompt", default=os.environ.get("SMOLVLA_PROMPT", "grab the block.")) parser.add_argument("--warmup", type=int, default=1) parser.add_argument("--loops", type=int, default=10) @@ -167,6 +178,7 @@ def main() -> int: width=args.width, height=args.height, state_dim=args.state_dim, + initial_noise_dim=args.initial_noise_dim, prompt=args.prompt, image_names=image_names, ) diff --git a/robot_server/test/test_server_latency.sh b/robot_server/test/test_server_latency.sh index 3073dff..de19516 100755 --- a/robot_server/test/test_server_latency.sh +++ b/robot_server/test/test_server_latency.sh @@ -7,9 +7,9 @@ set -e # bash robot_server/test/test_server_latency.sh # # Positional args: -# $1: model-type, e.g. smolvla / pi0 +# $1: model-type, e.g. smolvla / pi0 / starvla # $2: backend, e.g. mac-cpu / mac-metal / linux-cpu / linux-cuda -# $3: test-suite, e.g. smolvla-libero / smolvla-so101 / pi0-libero +# $3: test-suite, e.g. smolvla-libero / smolvla-so101 / pi0-libero / starvla-bridge ROBOT_CPP_ROOT="${ROBOT_CPP_ROOT:?ROBOT_CPP_ROOT must be set}" GGUF_DIR="${GGUF_DIR:?GGUF_DIR must be set}" MODEL_TYPE="${1:-${MODEL_TYPE:-smolvla}}" @@ -39,6 +39,17 @@ case "${MODEL_TYPE}" in ACTION_DECODER_GGUF="${ACTION_DECODER_GGUF:-${GGUF_DIR}/${MODEL_BASENAME}.action_decoder.gguf}" LLM_GGUF="${LLM_GGUF:-${GGUF_DIR}/${MODEL_BASENAME}.llm.gguf}" ;; + starvla) + mapfile -t LLM_CANDIDATES < <(find "${GGUF_DIR}" -maxdepth 1 -type f -name 'qwen-*.gguf' -print) + mapfile -t MMPROJ_CANDIDATES < <(find "${GGUF_DIR}" -maxdepth 1 -type f -name 'mmproj-*.gguf' -print) + mapfile -t POLICY_CANDIDATES < <(find "${GGUF_DIR}" -maxdepth 1 -type f -name '*policy*.gguf' -print) + [[ ${#LLM_CANDIDATES[@]} -eq 1 ]] || { echo "expected one Qwen GGUF in ${GGUF_DIR}" >&2; exit 1; } + [[ ${#MMPROJ_CANDIDATES[@]} -eq 1 ]] || { echo "expected one mmproj GGUF in ${GGUF_DIR}" >&2; exit 1; } + [[ ${#POLICY_CANDIDATES[@]} -eq 1 ]] || { echo "expected one policy GGUF in ${GGUF_DIR}" >&2; exit 1; } + LLM_GGUF="${LLM_GGUF:-${LLM_CANDIDATES[0]}}" + MMPROJ_GGUF="${MMPROJ_GGUF:-${MMPROJ_CANDIDATES[0]}}" + POLICY_GGUF="${POLICY_GGUF:-${POLICY_CANDIDATES[0]}}" + ;; *) echo "unsupported MODEL_TYPE=${MODEL_TYPE}" >&2 exit 1 @@ -99,6 +110,12 @@ case "${TEST_SUITE}" in IMAGE_HEIGHT="${IMAGE_HEIGHT:-256}" STATE_DIM="${STATE_DIM:-8}" ;; + starvla-bridge) + IMAGE_NAMES="${IMAGE_NAMES:-${IMAGE_NAME:-image_0}}" + IMAGE_WIDTH="${IMAGE_WIDTH:-224}" + IMAGE_HEIGHT="${IMAGE_HEIGHT:-224}" + STATE_DIM="${STATE_DIM:-0}" + ;; *) echo "unsupported TEST_SUITE=${TEST_SUITE}" >&2 exit 1 @@ -108,6 +125,7 @@ WARMUP="${WARMUP:-5}" LOOPS="${LOOPS:-100}" SERVER_WAIT_S="${SERVER_WAIT_S:-120}" DTYPE="${DTYPE:-f32}" +NOISE_SEED="${NOISE_SEED:--1}" PYTHON="${PYTHON:-python3}" # ==================================== @@ -155,12 +173,13 @@ run_latency_case() { TOKENIZER_GGUF="${TOKENIZER_GGUF:-}" \ STATE_GGUF="${STATE_GGUF:-}" \ ACTION_DECODER_GGUF="${ACTION_DECODER_GGUF:-}" \ + POLICY_GGUF="${POLICY_GGUF:-}" \ HOST="${HOST}" \ PORT="${PORT}" \ THREADS="${threads}" \ TASK="${PROMPT}" \ NOISE_MODE="gaussian" \ - NOISE_SEED="-1" \ + NOISE_SEED="${NOISE_SEED}" \ bash "${LAUNCH_SHELL}" "${MODEL_TYPE}" >"${server_log}" 2>&1 & SERVER_PID=$! diff --git a/src/model-cli.cpp b/src/model-cli.cpp index bc82d8e..8dc8441 100644 --- a/src/model-cli.cpp +++ b/src/model-cli.cpp @@ -1,7 +1,9 @@ // model-cli.cpp — common robotcpp::Model CLI frontend #include "models/model.h" +#include "models/argument_parse.h" #include "models/smolvla/smolvla_engine.h" +#include "llama.h" #include "stb_image.h" #include @@ -26,16 +28,11 @@ struct loaded_image { int stride_bytes = 0; }; -bool parse_model_type(const std::string & value, robotcpp::model_type & out) { - if (value == "smolvla") { - out = robotcpp::model_type::smolvla; - return true; - } - if (value == "pi0") { - out = robotcpp::model_type::pi0; - return true; +void quiet_llama_log_callback(ggml_log_level level, const char * text, void * user_data) { + (void)user_data; + if (level == GGML_LOG_LEVEL_ERROR) { + std::fputs(text, stderr); } - return false; } bool parse_noise_mode(const std::string & value, int & out_mode) { @@ -54,13 +51,14 @@ void print_usage(const char * prog) { std::fprintf(stderr, "\nModel CLI - robotcpp::Model frontend\n\n"); std::fprintf(stderr, "Usage:\n"); std::fprintf(stderr, " %s --model-type smolvla [options]\n", prog); - std::fprintf(stderr, " %s --model-type pi0 [options]\n\n", prog); + std::fprintf(stderr, " %s --model-type pi0 [options]\n", prog); + std::fprintf(stderr, " %s --model-type starvla --llm --mmproj --policy [options]\n\n", prog); std::fprintf(stderr, "Common options:\n"); - std::fprintf(stderr, " --model-type Model type (default: smolvla)\n"); + std::fprintf(stderr, " --model-type smolvla|pi0|starvla\n" + " (default: smolvla)\n"); std::fprintf(stderr, " --image Input image (repeatable; order matches --image-name)\n"); - std::fprintf( - stderr, - " --image-name Observation image name (repeatable; default: image for single-image input)\n"); + std::fprintf(stderr, + " --image-name Observation image name (default: image_0 for StarVLA, image otherwise)\n"); std::fprintf(stderr, " --state Proprio/state values (comma-separated)\n"); std::fprintf(stderr, " --task Task instruction (default: \"grab the block.\")\n"); std::fprintf(stderr, " --threads Number of threads (default: auto)\n"); @@ -82,6 +80,13 @@ void print_usage(const char * prog) { std::fprintf(stderr, " --tokenizer Tokenizer GGUF path\n"); std::fprintf(stderr, " --state-gguf State projector GGUF path\n"); std::fprintf(stderr, " --action-decoder Action decoder GGUF path\n"); + std::fprintf(stderr, "\nStarVLA options:\n"); + std::fprintf(stderr, " --policy StarVLA policy GGUF path (required)\n"); + std::fprintf(stderr, " --llm Qwen text GGUF path (required)\n"); + std::fprintf(stderr, " --mmproj Qwen vision GGUF path (required)\n"); + std::fprintf(stderr, " --n-batch Qwen batch size (default: 512)\n"); + std::fprintf(stderr, " --n-ctx Qwen context size (default: 2048)\n"); + std::fprintf(stderr, " --noise-seed GR00T/PI/PI_v3 noise seed, <0 means auto (default: -1)\n"); } bool parse_state(const char * csv, std::vector & out) { @@ -128,7 +133,7 @@ int main(int argc, char ** argv) { std::vector image_paths; std::vector image_names; std::string state_csv; - std::string task; + std::string task = "grab the block."; for (int i = 1; i < argc; i++) { std::string arg = argv[i]; @@ -139,10 +144,12 @@ int main(int argc, char ** argv) { } else if (arg == "-v" || arg == "--verbose") { args.verbosity++; } else if (arg == "--model-type" && i + 1 < argc) { - if (!parse_model_type(argv[++i], args.type)) { + if (!robotcpp::parse_model_type(argv[++i], args.type)) { std::fprintf(stderr, "Error: unsupported model type '%s'\n", argv[i]); return 1; } + } else if (arg == "--policy" && i + 1 < argc) { + args.policy_path = argv[++i]; } else if (arg == "--llm" && i + 1 < argc) { args.llm_path = argv[++i]; } else if (arg == "--mmproj" && i + 1 < argc) { @@ -168,18 +175,34 @@ int main(int argc, char ** argv) { } else if (arg == "--task" && i + 1 < argc) { task = argv[++i]; } else if (arg == "--threads" && i + 1 < argc) { - args.threads = std::atoi(argv[++i]); + const char * value = argv[++i]; + if (!robotcpp::parse_integer_argument(value, args.threads)) { + std::fprintf(stderr, "Error: invalid --threads value '%s'\n", value); + return 1; + } } else if (arg == "--n-batch" && i + 1 < argc) { - args.n_batch = std::atoi(argv[++i]); + const char * value = argv[++i]; + if (!robotcpp::parse_integer_argument(value, args.n_batch)) { + std::fprintf(stderr, "Error: invalid --n-batch value '%s'\n", value); + return 1; + } } else if (arg == "--n-ctx" && i + 1 < argc) { - args.n_ctx = std::atoi(argv[++i]); + const char * value = argv[++i]; + if (!robotcpp::parse_integer_argument(value, args.n_ctx)) { + std::fprintf(stderr, "Error: invalid --n-ctx value '%s'\n", value); + return 1; + } } else if (arg == "--noise-mode" && i + 1 < argc) { if (!parse_noise_mode(argv[++i], args.noise_mode)) { std::fprintf(stderr, "Error: invalid noise mode '%s'\n", argv[i]); return 1; } } else if (arg == "--noise-seed" && i + 1 < argc) { - args.noise_seed = std::atoll(argv[++i]); + const char * value = argv[++i]; + if (!robotcpp::parse_integer_argument(value, args.noise_seed)) { + std::fprintf(stderr, "Error: invalid --noise-seed value '%s'\n", value); + return 1; + } } else { std::fprintf(stderr, "Error: unknown argument '%s'\n", arg.c_str()); print_usage(argv[0]); @@ -187,23 +210,46 @@ int main(int argc, char ** argv) { } } + if (args.threads < 0 || args.n_batch <= 0 || args.n_ctx <= 0) { + std::fprintf(stderr, "Error: --threads must be non-negative and --n-batch/--n-ctx must be positive\n"); + return 1; + } + if (robotcpp::is_starvla_model_type(args.type)) { + if (args.llm_path.empty() || args.mmproj_path.empty() || args.policy_path.empty()) { + std::fprintf(stderr, "Error: %s requires --llm --mmproj --policy\n", robotcpp::model_type_name(args.type)); + return 1; + } + if (args.noise_mode != SMOLVLA_NOISE_MODE_GAUSSIAN) { + std::fprintf(stderr, "Error: StarVLA does not support --noise-mode debug-sin; use Gaussian noise\n"); + return 1; + } + } + if (image_paths.empty()) { std::fprintf(stderr, "Error: --image is required\n"); print_usage(argv[0]); return 1; } + if (robotcpp::is_starvla_model_type(args.type) && image_paths.size() != 1) { + std::fprintf(stderr, "Error: %s requires exactly one --image\n", robotcpp::model_type_name(args.type)); + return 1; + } if (image_names.empty()) { if (image_paths.size() != 1) { std::fprintf(stderr, "Error: multiple --image inputs require one --image-name per image\n"); return 1; } - image_names.push_back("image"); + image_names.push_back(robotcpp::is_starvla_model_type(args.type) ? "image_0" : "image"); } if (image_names.size() != image_paths.size()) { std::fprintf(stderr, "Error: --image count (%zu) must match --image-name count (%zu)\n", image_paths.size(), image_names.size()); return 1; } + if (robotcpp::is_starvla_model_type(args.type) && image_names[0] != "image_0") { + std::fprintf(stderr, "Error: %s image must be named 'image_0'\n", robotcpp::model_type_name(args.type)); + return 1; + } std::vector state_vec; if (!parse_state(state_csv.c_str(), state_vec)) { @@ -218,6 +264,8 @@ int main(int argc, char ** argv) { } } + llama_log_set(args.verbosity > 0 ? nullptr : quiet_llama_log_callback, nullptr); + const auto init_start = std::chrono::high_resolution_clock::now(); std::string error; std::unique_ptr model; diff --git a/src/models/argument_parse.h b/src/models/argument_parse.h new file mode 100644 index 0000000..0ad5b60 --- /dev/null +++ b/src/models/argument_parse.h @@ -0,0 +1,27 @@ +#pragma once + +#include +#include +#include +#include + +namespace robotcpp { + +template bool parse_integer_argument(const char * value, Integer & output) { + static_assert(std::is_integral::value && !std::is_same::value, + "Integer must be a non-bool integral type"); + if (value == nullptr || value[0] == '\0') { + return false; + } + + Integer parsed = 0; + const char * end = value + std::strlen(value); + const std::from_chars_result result = std::from_chars(value, end, parsed, 10); + if (result.ec != std::errc{} || result.ptr != end) { + return false; + } + output = parsed; + return true; +} + +} // namespace robotcpp diff --git a/src/models/ggml_backend.cpp b/src/models/ggml_backend.cpp index 062d895..65ef5ed 100644 --- a/src/models/ggml_backend.cpp +++ b/src/models/ggml_backend.cpp @@ -6,7 +6,7 @@ #include #include -static const char * backend_mode_name(backend_mode mode) { +const char * backend_mode_name(backend_mode mode) { switch (mode) { case backend_mode::cuda: return "cuda"; diff --git a/src/models/ggml_backend.h b/src/models/ggml_backend.h index 08b735d..a044f6f 100644 --- a/src/models/ggml_backend.h +++ b/src/models/ggml_backend.h @@ -11,6 +11,8 @@ enum class backend_mode { metal, }; +const char * backend_mode_name(backend_mode mode); + struct backend_buft_policy { ggml_backend_buffer_type_t model_buft = nullptr; ggml_backend_buffer_type_t runtime_buft = nullptr; diff --git a/src/models/model.h b/src/models/model.h index 587957c..d89b9d3 100644 --- a/src/models/model.h +++ b/src/models/model.h @@ -10,8 +10,13 @@ namespace robotcpp { enum class model_type { smolvla, pi0, + starvla, }; +const char * model_type_name(model_type type); +bool parse_model_type(const std::string & value, model_type & out); +bool is_starvla_model_type(model_type type); + struct model_image { std::string name; const uint8_t * data = nullptr; @@ -24,6 +29,7 @@ struct model_image { struct observation { std::vector images; std::vector state; + std::vector initial_noise; std::string task; }; @@ -59,6 +65,9 @@ struct model_args { std::string tokenizer_path; std::string state_path; std::string action_decoder_path; + + // starvla + std::string policy_path; }; class Model { diff --git a/src/models/model_factory.cpp b/src/models/model_factory.cpp index d6b1aa3..3f887ad 100644 --- a/src/models/model_factory.cpp +++ b/src/models/model_factory.cpp @@ -2,6 +2,9 @@ #include "models/pi0/pi0_model.h" #include "models/smolvla/smolvla_model.h" +#ifdef ROBOT_CPP_BUILD_STARVLA +#include "models/starvla/starvla_model.h" +#endif namespace robotcpp { @@ -13,8 +16,16 @@ bool make_model(const model_args & args, std::unique_ptr & out, std::stri if (args.type == model_type::pi0) { return make_pi0_model(args, out, error); } + if (args.type == model_type::starvla) { +#ifdef ROBOT_CPP_BUILD_STARVLA + return make_starvla_model(args, out, error); +#else + error = "StarVLA support was not built; configure with -DROBOT_CPP_BUILD_STARVLA=ON"; + return false; +#endif + } - error = "unsupported model type"; + error = std::string("unsupported model type: ") + model_type_name(args.type); return false; } diff --git a/src/models/model_type.cpp b/src/models/model_type.cpp new file mode 100644 index 0000000..6a4d3d5 --- /dev/null +++ b/src/models/model_type.cpp @@ -0,0 +1,49 @@ +#include "models/model.h" + +#include + +namespace robotcpp { +namespace { + +struct model_type_entry { + model_type type; + const char * name; +}; + +constexpr std::array MODEL_TYPES = {{ + {model_type::smolvla, "smolvla"}, + {model_type::pi0, "pi0"}, + {model_type::starvla, "starvla"}, +}}; + +const model_type_entry * find_entry(model_type type) { + for (const model_type_entry & entry : MODEL_TYPES) { + if (entry.type == type) { + return &entry; + } + } + return nullptr; +} + +} // namespace + +const char * model_type_name(model_type type) { + const model_type_entry * entry = find_entry(type); + return entry ? entry->name : "unknown"; +} + +bool parse_model_type(const std::string & value, model_type & out) { + for (const model_type_entry & entry : MODEL_TYPES) { + if (value == entry.name) { + out = entry.type; + return true; + } + } + return false; +} + +bool is_starvla_model_type(model_type type) { + return type == model_type::starvla; +} + +} // namespace robotcpp diff --git a/src/models/pi0/pi0_model.cpp b/src/models/pi0/pi0_model.cpp index 5d62866..4d46a0d 100644 --- a/src/models/pi0/pi0_model.cpp +++ b/src/models/pi0/pi0_model.cpp @@ -80,6 +80,10 @@ bool Pi0Model::predict(const observation & obs, model_result & out, std::string error = "Pi0 model is not initialized"; return false; } + if (!obs.initial_noise.empty()) { + error = "Pi0 does not accept explicit initial noise"; + return false; + } if (obs.images.empty()) { error = "Pi0 requires at least one image"; return false; diff --git a/src/models/smolvla/smolvla_model.cpp b/src/models/smolvla/smolvla_model.cpp index a718fd2..db6567c 100644 --- a/src/models/smolvla/smolvla_model.cpp +++ b/src/models/smolvla/smolvla_model.cpp @@ -72,6 +72,10 @@ bool SmolVLAModel::predict(const observation & obs, model_result & out, std::str error = "SmolVLA model is not initialized"; return false; } + if (!obs.initial_noise.empty()) { + error = "SmolVLA does not accept explicit initial noise"; + return false; + } if (obs.images.empty()) { error = "SmolVLA requires at least one image"; return false; diff --git a/src/models/starvla/fast_codec.cpp b/src/models/starvla/fast_codec.cpp index 3a01ac9..9ce67fd 100644 --- a/src/models/starvla/fast_codec.cpp +++ b/src/models/starvla/fast_codec.cpp @@ -1,14 +1,8 @@ #include "models/starvla/fast_codec.h" -#include "nlohmann/json.hpp" - #include -#include #include -#include -#include #include -#include #include #include #include @@ -16,87 +10,34 @@ namespace robotcpp::starvla { namespace { -using Json = nlohmann::json; - -constexpr size_t kMaximumJsonBytes = 16U * 1024U * 1024U; -constexpr size_t kOfficialVocabSize = 2048U; -constexpr size_t kMaximumVocabSize = 65536U; -constexpr size_t kMaximumTimeHorizon = 1024U; -constexpr size_t kMaximumActionDim = 1024U; -constexpr size_t kMaximumBatchSize = 1024U; -constexpr size_t kMaximumTokenSequence = 4096U; -constexpr size_t kMaximumGeneratedSequence = 2048U; -constexpr size_t kMaximumDecodedBytes = 1024U * 1024U; -constexpr size_t kMaximumOutputScalars = 16U * 1024U * 1024U; +constexpr size_t kMaximumVocabSize = 65536U; +constexpr size_t kMaximumTimeHorizon = 1024U; +constexpr size_t kMaximumActionDim = 1024U; +constexpr size_t kMaximumBatchSize = 1024U; +constexpr size_t kMaximumTokenSequence = 4096U; +constexpr size_t kMaximumGeneratedSequence = 2048U; +constexpr size_t kMaximumDecodedBytes = 1024U * 1024U; +constexpr size_t kMaximumOutputScalars = 16U * 1024U * 1024U; constexpr uint64_t kMaximumIdctMultiplyAdds = 64ULL * 1024ULL * 1024ULL; -constexpr const char * kActionTokenPrefix = " kMaximumJsonBytes) { - error = "StarVLA FAST JSON asset exceeds the 16 MiB limit: " + path.string(); - return false; - } - - std::ifstream stream(path, std::ios::binary); - if (!stream) { - error = "cannot open StarVLA FAST JSON asset: " + path.string(); - return false; - } - std::string contents; - contents.reserve(static_cast(size)); - std::array buffer{}; - while (stream) { - stream.read(buffer.data(), static_cast(buffer.size())); - const std::streamsize count = stream.gcount(); - if (count <= 0) { - continue; - } - const size_t chunk_size = static_cast(count); - if (contents.size() > kMaximumJsonBytes - chunk_size) { - error = "StarVLA FAST JSON asset exceeds the 16 MiB limit while reading: " + - path.string(); - return false; - } - contents.append(buffer.data(), chunk_size); - } - if (!stream.eof() || stream.bad()) { - error = "cannot read StarVLA FAST JSON asset: " + path.string(); - return false; - } - output = Json::parse(contents, nullptr, false); - if (output.is_discarded()) { - error = "cannot parse StarVLA FAST JSON asset: " + path.string(); - return false; - } - return true; -} +constexpr double kPi = 3.141592653589793238462643383279502884; -bool decode_utf8_strict(const std::string & input, std::vector & output, - std::string & error) { +bool decode_utf8_strict(const std::string & input, std::vector & output, std::string & error) { output.clear(); for (size_t i = 0; i < input.size();) { const uint8_t first = static_cast(input[i]); - uint32_t value = 0; - size_t length = 0; + uint32_t value = 0; + size_t length = 0; if (first <= 0x7fU) { - value = first; + value = first; length = 1; } else if (first >= 0xc2U && first <= 0xdfU) { - value = first & 0x1fU; + value = first & 0x1fU; length = 2; } else if (first >= 0xe0U && first <= 0xefU) { - value = first & 0x0fU; + value = first & 0x0fU; length = 3; } else if (first >= 0xf0U && first <= 0xf4U) { - value = first & 0x07U; + value = first & 0x07U; length = 4; } else { error = "StarVLA FAST tokenizer vocabulary contains invalid UTF-8"; @@ -114,9 +55,8 @@ bool decode_utf8_strict(const std::string & input, std::vector & outpu } value = (value << 6U) | (continuation & 0x3fU); } - const bool overlong = (length == 2 && value < 0x80U) || - (length == 3 && value < 0x800U) || - (length == 4 && value < 0x10000U); + const bool overlong = + (length == 2 && value < 0x80U) || (length == 3 && value < 0x800U) || (length == 4 && value < 0x10000U); if (overlong || value > 0x10ffffU || (value >= 0xd800U && value <= 0xdfffU)) { error = "StarVLA FAST tokenizer vocabulary contains a non-scalar UTF-8 value"; return false; @@ -152,16 +92,14 @@ std::unordered_map byte_level_inverse_alphabet() { return result; } -bool compile_token_bytes(const std::vector & vocab_by_id, - std::vector> & token_bytes, +bool compile_token_bytes(const std::vector & vocab_by_id, std::vector> & token_bytes, std::string & error) { const auto inverse_alphabet = byte_level_inverse_alphabet(); token_bytes.clear(); token_bytes.reserve(vocab_by_id.size()); for (size_t token_id = 0; token_id < vocab_by_id.size(); ++token_id) { if (vocab_by_id[token_id].empty()) { - error = "StarVLA FAST tokenizer has an empty vocabulary piece at ID " + - std::to_string(token_id); + error = "StarVLA FAST tokenizer has an empty vocabulary piece at ID " + std::to_string(token_id); return false; } std::vector piece_codepoints; @@ -196,17 +134,17 @@ void decode_utf8_lossy(const std::vector & input, std::vector continue; } - size_t length = 0; + size_t length = 0; uint32_t value = 0; if (first >= 0xc2U && first <= 0xdfU) { length = 2; - value = first & 0x1fU; + value = first & 0x1fU; } else if (first >= 0xe0U && first <= 0xefU) { length = 3; - value = first & 0x0fU; + value = first & 0x0fU; } else if (first >= 0xf0U && first <= 0xf4U) { length = 4; - value = first & 0x07U; + value = first & 0x07U; } else { output.push_back(0xfffdU); ++i; @@ -217,13 +155,11 @@ void decode_utf8_lossy(const std::vector & input, std::vector output.push_back(0xfffdU); break; } - const uint8_t second = input[i + 1]; + const uint8_t second = input[i + 1]; const bool second_is_continuation = (second & 0xc0U) == 0x80U; - const bool second_in_scalar_range = - !(first == 0xe0U && second < 0xa0U) && - !(first == 0xedU && second > 0x9fU) && - !(first == 0xf0U && second < 0x90U) && - !(first == 0xf4U && second > 0x8fU); + const bool second_in_scalar_range = !(first == 0xe0U && second < 0xa0U) && + !(first == 0xedU && second > 0x9fU) && + !(first == 0xf0U && second < 0x90U) && !(first == 0xf4U && second > 0x8fU); if (!second_is_continuation || !second_in_scalar_range) { output.push_back(0xfffdU); ++i; @@ -231,12 +167,12 @@ void decode_utf8_lossy(const std::vector & input, std::vector } value = (value << 6U) | (second & 0x3fU); - bool invalid = false; + bool invalid = false; size_t consumed_prefix = 2; for (size_t j = 2; j < length; ++j) { if (i + j >= input.size()) { output.push_back(0xfffdU); - i = input.size(); + i = input.size(); invalid = true; break; } @@ -258,202 +194,8 @@ void decode_utf8_lossy(const std::vector & input, std::vector } } -bool parse_positive_size(const Json & value, const char * name, size_t & output, - std::string & error) { - if (!value.is_number_integer()) { - error = std::string("StarVLA FAST ") + name + " must be an integer"; - return false; - } - try { - const int64_t parsed = value.get(); - if (parsed <= 0 || static_cast(parsed) > - static_cast(std::numeric_limits::max())) { - error = std::string("StarVLA FAST ") + name + " is out of range"; - return false; - } - output = static_cast(parsed); - return true; - } catch (const std::exception &) { - error = std::string("StarVLA FAST ") + name + " is out of range"; - return false; - } -} - -bool parse_processor_config(const Json & json, size_t time_horizon_override, - size_t action_dim_override, FastCodecConfig & config, - std::string & error) { - if (!json.is_object() || !json.contains("processor_class") || - json["processor_class"] != "UniversalActionProcessor" || !json.contains("scale") || - !json["scale"].is_number() || !json.contains("vocab_size") || - !json.contains("min_token") || !json["min_token"].is_number_integer()) { - error = "StarVLA FAST processor_config.json has an incompatible schema"; - return false; - } - config.scale = json["scale"].get(); - if (!std::isfinite(config.scale) || config.scale == 0.0) { - error = "StarVLA FAST processor scale must be finite and non-zero"; - return false; - } - if (!parse_positive_size(json["vocab_size"], "vocab_size", config.vocab_size, error)) { - return false; - } - if (config.vocab_size != kOfficialVocabSize) { - error = "StarVLA FAST pinned vocabulary must contain exactly 2048 tokens"; - return false; - } - try { - if (json["min_token"].is_number_unsigned()) { - const uint64_t min_token = json["min_token"].get(); - if (min_token > static_cast(std::numeric_limits::max())) { - error = "StarVLA FAST min_token is out of int32 range"; - return false; - } - config.min_token = static_cast(min_token); - } else { - const int64_t min_token = json["min_token"].get(); - if (min_token < std::numeric_limits::min() || - min_token > std::numeric_limits::max()) { - error = "StarVLA FAST min_token is out of int32 range"; - return false; - } - config.min_token = static_cast(min_token); - } - } catch (const std::exception &) { - error = "StarVLA FAST min_token is out of int32 range"; - return false; - } - - auto choose_dimension = [&](const char * name, size_t override_value, size_t & target) { - if (override_value != 0) { - target = override_value; - return true; - } - if (!json.contains(name) || json[name].is_null()) { - error = std::string("StarVLA FAST ") + name + - " is absent; pass the policy dimension explicitly"; - return false; - } - return parse_positive_size(json[name], name, target, error); - }; - return choose_dimension("time_horizon", time_horizon_override, config.time_horizon) && - choose_dimension("action_dim", action_dim_override, config.action_dim); -} - -bool parse_tokenizer_vocab(const Json & json, size_t expected_vocab_size, - std::vector & vocab_by_id, std::string & error) { - if (!json.is_object() || !json.contains("version") || json["version"] != "1.0" || - !json.contains("added_tokens") || !json["added_tokens"].is_array() || - !json["added_tokens"].empty() || !json.contains("decoder") || - !json["decoder"].is_object() || !json["decoder"].contains("type") || - json["decoder"]["type"] != "ByteLevel" || !json.contains("model") || - !json["model"].is_object() || !json["model"].contains("type") || - json["model"]["type"] != "BPE" || !json["model"].contains("vocab") || - !json["model"]["vocab"].is_object()) { - error = "StarVLA FAST tokenizer.json is not the required ByteLevel BPE schema"; - return false; - } - const Json & decoder = json["decoder"]; - if (decoder.size() != 4 || !decoder.contains("add_prefix_space") || - decoder["add_prefix_space"] != true || !decoder.contains("trim_offsets") || - decoder["trim_offsets"] != true || !decoder.contains("use_regex") || - decoder["use_regex"] != true) { - error = "StarVLA FAST tokenizer.json has an incompatible ByteLevel decoder contract"; - return false; - } - const Json & vocab = json["model"]["vocab"]; - if (vocab.size() != expected_vocab_size) { - error = "StarVLA FAST tokenizer vocabulary size does not match processor_config.json"; - return false; - } - vocab_by_id.assign(expected_vocab_size, std::string()); - std::vector seen(expected_vocab_size, false); - for (auto iterator = vocab.begin(); iterator != vocab.end(); ++iterator) { - if (!iterator.value().is_number_integer()) { - error = "StarVLA FAST tokenizer vocabulary ID is not an integer"; - return false; - } - int64_t token_id = -1; - try { - token_id = iterator.value().get(); - } catch (const std::exception &) { - error = "StarVLA FAST tokenizer vocabulary ID is out of range"; - return false; - } - if (token_id < 0 || static_cast(token_id) >= expected_vocab_size || - seen[static_cast(token_id)]) { - error = "StarVLA FAST tokenizer vocabulary IDs are not a bijection"; - return false; - } - seen[static_cast(token_id)] = true; - vocab_by_id[static_cast(token_id)] = iterator.key(); - } - return true; -} - -bool parse_action_index(const std::string & value, size_t & index) { - const std::string prefix(kActionTokenPrefix); - if (value.size() <= prefix.size() + 1 || value.compare(0, prefix.size(), prefix) != 0 || - value.back() != '>') { - return false; - } - const std::string digits = value.substr(prefix.size(), value.size() - prefix.size() - 1); - if (digits.empty() || (digits.size() > 1 && digits.front() == '0')) { - return false; - } - size_t parsed = 0; - for (char character : digits) { - if (character < '0' || character > '9') { - return false; - } - const size_t digit = static_cast(character - '0'); - if (parsed > (std::numeric_limits::max() - digit) / 10U) { - return false; - } - parsed = parsed * 10U + digit; - } - index = parsed; - return true; -} - -bool parse_action_map(const Json & json, size_t vocab_size, - std::vector & fast_to_vlm, std::string & error) { - if (!json.is_object() || json.size() != vocab_size) { - error = "StarVLA FAST action-token map must contain exactly one entry per FAST token"; - return false; - } - fast_to_vlm.assign(vocab_size, -1); - std::vector seen(vocab_size, false); - std::unordered_set vlm_ids; - for (auto iterator = json.begin(); iterator != json.end(); ++iterator) { - size_t fast_id = 0; - if (!parse_action_index(iterator.key(), fast_id) || fast_id >= vocab_size || seen[fast_id]) { - error = "StarVLA FAST action-token map has a malformed or duplicate token name"; - return false; - } - if (!iterator.value().is_number_integer()) { - error = "StarVLA FAST action-token map contains a non-integer VLM ID"; - return false; - } - int64_t vlm_id = -1; - try { - vlm_id = iterator.value().get(); - } catch (const std::exception &) { - error = "StarVLA FAST action-token VLM ID is out of range"; - return false; - } - if (vlm_id < 0 || vlm_id > std::numeric_limits::max() || - !vlm_ids.insert(static_cast(vlm_id)).second) { - error = "StarVLA FAST action-token VLM IDs must be unique non-negative int32 values"; - return false; - } - seen[fast_id] = true; - fast_to_vlm[fast_id] = static_cast(vlm_id); - } - return true; -} - -bool checked_action_count(const FastCodecConfig & config, size_t batch_size, - size_t & per_sample, size_t & total, std::string & error) { +bool checked_action_count(const FastCodecConfig & config, size_t batch_size, size_t & per_sample, size_t & total, + std::string & error) { if (config.vocab_size > kMaximumVocabSize || config.time_horizon > kMaximumTimeHorizon || config.action_dim > kMaximumActionDim) { error = "StarVLA FAST codec dimensions exceed the runtime safety limits"; @@ -477,9 +219,9 @@ bool checked_action_count(const FastCodecConfig & config, size_t batch_size, error = "StarVLA FAST output tensor exceeds the runtime scalar limit"; return false; } - const uint64_t horizon = static_cast(config.time_horizon); + const uint64_t horizon = static_cast(config.time_horizon); const uint64_t action_dim = static_cast(config.action_dim); - const uint64_t batch = static_cast(batch_size); + const uint64_t batch = static_cast(batch_size); if (horizon > kMaximumIdctMultiplyAdds / horizon) { error = "StarVLA FAST inverse DCT exceeds the runtime work limit"; return false; @@ -501,8 +243,7 @@ bool checked_action_count(const FastCodecConfig & config, size_t batch_size, FastCodec::FastCodec(FastCodecConfig config, std::vector> token_bytes, std::vector fast_to_vlm_id) - : config_(config), token_bytes_(std::move(token_bytes)), - fast_to_vlm_id_(std::move(fast_to_vlm_id)) { + : config_(config), token_bytes_(std::move(token_bytes)), fast_to_vlm_id_(std::move(fast_to_vlm_id)) { vlm_to_fast_id_.reserve(fast_to_vlm_id_.size()); for (size_t fast_id = 0; fast_id < fast_to_vlm_id_.size(); ++fast_id) { vlm_to_fast_id_.emplace_back(fast_to_vlm_id_[fast_id], static_cast(fast_id)); @@ -510,13 +251,11 @@ FastCodec::FastCodec(FastCodecConfig config, std::vector> t std::sort(vlm_to_fast_id_.begin(), vlm_to_fast_id_.end()); } -std::unique_ptr FastCodec::create(FastCodecConfig config, - std::vector vocab_by_id, - std::vector fast_to_vlm_id, - std::string & error) { +std::unique_ptr FastCodec::create(FastCodecConfig config, std::vector vocab_by_id, + std::vector fast_to_vlm_id, std::string & error) { error.clear(); - if (!std::isfinite(config.scale) || config.scale == 0.0 || config.vocab_size == 0 || - config.time_horizon == 0 || config.action_dim == 0) { + if (!std::isfinite(config.scale) || config.scale == 0.0 || config.vocab_size == 0 || config.time_horizon == 0 || + config.action_dim == 0) { error = "StarVLA FAST codec dimensions and scale must be non-zero and finite"; return nullptr; } @@ -524,13 +263,12 @@ std::unique_ptr FastCodec::create(FastCodecConfig config, error = "StarVLA FAST vocabulary exceeds the int32 token-ID range"; return nullptr; } - if (vocab_by_id.size() != config.vocab_size || - fast_to_vlm_id.size() != config.vocab_size) { + if (vocab_by_id.size() != config.vocab_size || fast_to_vlm_id.size() != config.vocab_size) { error = "StarVLA FAST codec vocabulary or action-token map has the wrong size"; return nullptr; } size_t per_sample = 0; - size_t total = 0; + size_t total = 0; if (!checked_action_count(config, 1, per_sample, total, error)) { return nullptr; } @@ -545,34 +283,26 @@ std::unique_ptr FastCodec::create(FastCodecConfig config, if (!compile_token_bytes(vocab_by_id, token_bytes, error)) { return nullptr; } - return std::unique_ptr( - new FastCodec(config, std::move(token_bytes), std::move(fast_to_vlm_id))); + return std::unique_ptr(new FastCodec(config, std::move(token_bytes), std::move(fast_to_vlm_id))); } -std::unique_ptr FastCodec::create_compiled( - FastCodecConfig config, std::vector token_offsets, - std::vector token_bytes, std::vector fast_to_vlm_id, - std::string & error) { +std::unique_ptr FastCodec::create_compiled(FastCodecConfig config, std::vector token_offsets, + std::vector token_bytes, + std::vector fast_to_vlm_id, std::string & error) { error.clear(); - if (!std::isfinite(config.scale) || config.scale == 0.0 || - config.vocab_size == 0 || config.time_horizon == 0 || - config.action_dim == 0 || - config.vocab_size > - static_cast(std::numeric_limits::max())) { + if (!std::isfinite(config.scale) || config.scale == 0.0 || config.vocab_size == 0 || config.time_horizon == 0 || + config.action_dim == 0 || config.vocab_size > static_cast(std::numeric_limits::max())) { error = "StarVLA FAST compiled codec dimensions and scale are invalid"; return nullptr; } - if (config.vocab_size == std::numeric_limits::max() || - token_offsets.size() != config.vocab_size + 1U || - fast_to_vlm_id.size() != config.vocab_size || - token_offsets.empty() || token_offsets.front() != 0 || - token_offsets.back() < 0 || - static_cast(token_offsets.back()) != token_bytes.size()) { + if (config.vocab_size == std::numeric_limits::max() || token_offsets.size() != config.vocab_size + 1U || + fast_to_vlm_id.size() != config.vocab_size || token_offsets.empty() || token_offsets.front() != 0 || + token_offsets.back() < 0 || static_cast(token_offsets.back()) != token_bytes.size()) { error = "StarVLA FAST compiled codec tensor shapes are incompatible"; return nullptr; } size_t per_sample = 0; - size_t total = 0; + size_t total = 0; if (!checked_action_count(config, 1, per_sample, total, error)) { return nullptr; } @@ -580,8 +310,7 @@ std::unique_ptr FastCodec::create_compiled( std::unordered_set unique_vlm_ids; for (int32_t vlm_id : fast_to_vlm_id) { if (vlm_id < 0 || !unique_vlm_ids.insert(vlm_id).second) { - error = - "StarVLA FAST compiled action-token IDs must be unique and non-negative"; + error = "StarVLA FAST compiled action-token IDs must be unique and non-negative"; return nullptr; } } @@ -590,45 +319,14 @@ std::unique_ptr FastCodec::create_compiled( pieces.reserve(config.vocab_size); for (size_t index = 0; index < config.vocab_size; ++index) { const int32_t begin = token_offsets[index]; - const int32_t end = token_offsets[index + 1U]; - if (begin < 0 || end <= begin || - static_cast(end) > token_bytes.size()) { + const int32_t end = token_offsets[index + 1U]; + if (begin < 0 || end <= begin || static_cast(end) > token_bytes.size()) { error = "StarVLA FAST compiled codec offsets are not strictly increasing"; return nullptr; } pieces.emplace_back(token_bytes.begin() + begin, token_bytes.begin() + end); } - return std::unique_ptr( - new FastCodec(config, std::move(pieces), std::move(fast_to_vlm_id))); -} - -std::unique_ptr FastCodec::load_hf_assets( - const std::filesystem::path & tokenizer_json, - const std::filesystem::path & processor_config_json, - const std::filesystem::path & action_token_map_json, - size_t time_horizon, size_t action_dim, std::string & error) { - error.clear(); - Json processor; - Json tokenizer; - Json action_map; - if (!read_json(processor_config_json, processor, error) || - !read_json(tokenizer_json, tokenizer, error) || - !read_json(action_token_map_json, action_map, error)) { - return nullptr; - } - - FastCodecConfig config; - std::vector vocab_by_id; - std::vector fast_to_vlm; - size_t per_sample = 0; - size_t total = 0; - if (!parse_processor_config(processor, time_horizon, action_dim, config, error) || - !checked_action_count(config, 1, per_sample, total, error) || - !parse_tokenizer_vocab(tokenizer, config.vocab_size, vocab_by_id, error) || - !parse_action_map(action_map, config.vocab_size, fast_to_vlm, error)) { - return nullptr; - } - return create(config, std::move(vocab_by_id), std::move(fast_to_vlm), error); + return std::unique_ptr(new FastCodec(config, std::move(pieces), std::move(fast_to_vlm_id))); } const FastCodecConfig & FastCodec::config() const { @@ -639,8 +337,8 @@ const std::vector & FastCodec::fast_to_vlm_ids() const { return fast_to_vlm_id_; } -bool FastCodec::map_fast_to_vlm(const std::vector & fast_ids, - std::vector & vlm_ids, std::string & error) const { +bool FastCodec::map_fast_to_vlm(const std::vector & fast_ids, std::vector & vlm_ids, + std::string & error) const { vlm_ids.clear(); error.clear(); if (fast_ids.size() > kMaximumTokenSequence) { @@ -659,8 +357,8 @@ bool FastCodec::map_fast_to_vlm(const std::vector & fast_ids, return true; } -bool FastCodec::map_vlm_to_fast(const std::vector & vlm_ids, - std::vector & fast_ids, std::string & error) const { +bool FastCodec::map_vlm_to_fast(const std::vector & vlm_ids, std::vector & fast_ids, + std::string & error) const { fast_ids.clear(); error.clear(); if (vlm_ids.size() > kMaximumTokenSequence) { @@ -671,9 +369,7 @@ bool FastCodec::map_vlm_to_fast(const std::vector & vlm_ids, for (int32_t vlm_id : vlm_ids) { const auto found = std::lower_bound( vlm_to_fast_id_.begin(), vlm_to_fast_id_.end(), vlm_id, - [](const std::pair & entry, int32_t value) { - return entry.first < value; - }); + [](const std::pair & entry, int32_t value) { return entry.first < value; }); if (found == vlm_to_fast_id_.end() || found->first != vlm_id) { error = "Qwen token ID is not present in the StarVLA FAST action-token map"; fast_ids.clear(); @@ -684,8 +380,7 @@ bool FastCodec::map_vlm_to_fast(const std::vector & vlm_ids, return true; } -bool FastCodec::extract_fast_tokens(const std::vector & generated_ids, - std::vector & fast_ids, +bool FastCodec::extract_fast_tokens(const std::vector & generated_ids, std::vector & fast_ids, std::string & error) const { fast_ids.clear(); error.clear(); @@ -696,9 +391,7 @@ bool FastCodec::extract_fast_tokens(const std::vector & generated_ids, for (int32_t vlm_id : generated_ids) { const auto found = std::lower_bound( vlm_to_fast_id_.begin(), vlm_to_fast_id_.end(), vlm_id, - [](const std::pair & entry, int32_t value) { - return entry.first < value; - }); + [](const std::pair & entry, int32_t value) { return entry.first < value; }); if (found != vlm_to_fast_id_.end() && found->first == vlm_id) { fast_ids.push_back(found->second); } @@ -706,8 +399,7 @@ bool FastCodec::extract_fast_tokens(const std::vector & generated_ids, return true; } -bool FastCodec::byte_level_decode(const std::vector & fast_ids, - std::vector & codepoints, +bool FastCodec::byte_level_decode(const std::vector & fast_ids, std::vector & codepoints, std::string & error) const { codepoints.clear(); error.clear(); @@ -742,8 +434,8 @@ bool FastCodec::byte_level_decode(const std::vector & fast_ids, return true; } -bool FastCodec::decode_fast_tokens(const std::vector> & batch_fast_ids, - FastDecodeResult & result, std::string & error) const { +bool FastCodec::decode_fast_tokens(const std::vector> & batch_fast_ids, FastDecodeResult & result, + std::string & error) const { result = {}; error.clear(); if (batch_fast_ids.empty()) { @@ -761,14 +453,14 @@ bool FastCodec::decode_fast_tokens(const std::vector> & bat } } size_t per_sample = 0; - size_t total = 0; + size_t total = 0; if (!checked_action_count(config_, batch_fast_ids.size(), per_sample, total, error)) { return false; } - result.batch_size = batch_fast_ids.size(); + result.batch_size = batch_fast_ids.size(); result.time_horizon = config_.time_horizon; - result.action_dim = config_.action_dim; + result.action_dim = config_.action_dim; result.actions.assign(total, 0.0); const double dc_scale = 1.0 / std::sqrt(static_cast(config_.time_horizon)); @@ -776,27 +468,22 @@ bool FastCodec::decode_fast_tokens(const std::vector> & bat for (size_t batch = 0; batch < batch_fast_ids.size(); ++batch) { std::vector codepoints; std::string sequence_error; - if (!byte_level_decode(batch_fast_ids[batch], codepoints, sequence_error) || - codepoints.size() != per_sample) { - error = "StarVLA FAST sequence " + std::to_string(batch) + ": " + - (sequence_error.empty() ? "decoded DCT coefficient shape mismatch" - : sequence_error); + if (!byte_level_decode(batch_fast_ids[batch], codepoints, sequence_error) || codepoints.size() != per_sample) { + error = "StarVLA FAST sequence " + std::to_string(batch) + ": " + + (sequence_error.empty() ? "decoded DCT coefficient shape mismatch" : sequence_error); result = {}; return false; } for (size_t action = 0; action < config_.action_dim; ++action) { - const double dc = - (static_cast(codepoints[action]) + config_.min_token) / config_.scale; + const double dc = (static_cast(codepoints[action]) + config_.min_token) / config_.scale; for (size_t time = 0; time < config_.time_horizon; ++time) { double value = dc_scale * dc; for (size_t frequency = 1; frequency < config_.time_horizon; ++frequency) { const size_t coefficient_index = frequency * config_.action_dim + action; const double coefficient = - (static_cast(codepoints[coefficient_index]) + config_.min_token) / - config_.scale; - const double angle = kPi * static_cast(frequency) * - static_cast(2U * time + 1U) / + (static_cast(codepoints[coefficient_index]) + config_.min_token) / config_.scale; + const double angle = kPi * static_cast(frequency) * static_cast(2U * time + 1U) / (2.0 * static_cast(config_.time_horizon)); value += ac_scale * coefficient * std::cos(angle); } @@ -807,12 +494,11 @@ bool FastCodec::decode_fast_tokens(const std::vector> & bat return true; } -bool FastCodec::decode_vlm_action_tokens( - const std::vector> & batch_vlm_ids, FastDecodeResult & result, - std::string & error) const { +bool FastCodec::decode_vlm_action_tokens(const std::vector> & batch_vlm_ids, + FastDecodeResult & result, std::string & error) const { if (batch_vlm_ids.empty() || batch_vlm_ids.size() > kMaximumBatchSize) { result = {}; - error = "StarVLA FAST action-token batch is empty or exceeds the runtime size limit"; + error = "StarVLA FAST action-token batch is empty or exceeds the runtime size limit"; return false; } std::vector> batch_fast_ids; @@ -828,12 +514,11 @@ bool FastCodec::decode_vlm_action_tokens( return decode_fast_tokens(batch_fast_ids, result, error); } -bool FastCodec::decode_generated_tokens( - const std::vector> & batch_generated_ids, FastDecodeResult & result, - std::string & error) const { +bool FastCodec::decode_generated_tokens(const std::vector> & batch_generated_ids, + FastDecodeResult & result, std::string & error) const { if (batch_generated_ids.empty() || batch_generated_ids.size() > kMaximumBatchSize) { result = {}; - error = "StarVLA FAST generated-token batch is empty or exceeds the runtime size limit"; + error = "StarVLA FAST generated-token batch is empty or exceeds the runtime size limit"; return false; } std::vector> batch_fast_ids; diff --git a/src/models/starvla/fast_codec.h b/src/models/starvla/fast_codec.h index b84314b..4c3742b 100644 --- a/src/models/starvla/fast_codec.h +++ b/src/models/starvla/fast_codec.h @@ -2,7 +2,6 @@ #include #include -#include #include #include #include @@ -11,76 +10,66 @@ namespace robotcpp::starvla { struct FastCodecConfig { - double scale = 0.0; - int32_t min_token = 0; - size_t vocab_size = 0; + double scale = 0.0; + int32_t min_token = 0; + size_t vocab_size = 0; size_t time_horizon = 0; - size_t action_dim = 0; + size_t action_dim = 0; }; struct FastDecodeResult { - size_t batch_size = 0; + size_t batch_size = 0; size_t time_horizon = 0; - size_t action_dim = 0; + size_t action_dim = 0; std::vector actions; }; class FastCodec { -public: - static std::unique_ptr create( - FastCodecConfig config, std::vector vocab_by_id, - std::vector fast_to_vlm_id, std::string & error); + public: + static std::unique_ptr create(FastCodecConfig config, std::vector vocab_by_id, + std::vector fast_to_vlm_id, std::string & error); // Constructs directly from the converter-compiled ByteLevel pieces stored // in policy GGUF. offsets has vocab_size + 1 entries and indexes the flat // byte buffer; no external tokenizer JSON is consulted. - static std::unique_ptr create_compiled( - FastCodecConfig config, std::vector token_offsets, - std::vector token_bytes, - std::vector fast_to_vlm_id, std::string & error); - - static std::unique_ptr load_hf_assets( - const std::filesystem::path & tokenizer_json, - const std::filesystem::path & processor_config_json, - const std::filesystem::path & action_token_map_json, - size_t time_horizon, size_t action_dim, std::string & error); + static std::unique_ptr create_compiled(FastCodecConfig config, std::vector token_offsets, + std::vector token_bytes, + std::vector fast_to_vlm_id, std::string & error); const FastCodecConfig & config() const; const std::vector & fast_to_vlm_ids() const; - bool map_fast_to_vlm(const std::vector & fast_ids, - std::vector & vlm_ids, std::string & error) const; - bool map_vlm_to_fast(const std::vector & vlm_ids, - std::vector & fast_ids, std::string & error) const; + bool map_fast_to_vlm(const std::vector & fast_ids, std::vector & vlm_ids, + std::string & error) const; + bool map_vlm_to_fast(const std::vector & vlm_ids, std::vector & fast_ids, + std::string & error) const; // Extracts every mapped action token from a generated Qwen sequence in order. // EOS stopping remains the generator's responsibility; ordinary EOS/pad/text // IDs in the returned sequence are ignored and do not terminate this scan. - bool extract_fast_tokens(const std::vector & generated_ids, - std::vector & fast_ids, std::string & error) const; - - // Exposed for focused parity diagnostics. This is the Hugging Face ByteLevel - // decoder output before min_token adjustment and inverse DCT. - bool byte_level_decode(const std::vector & fast_ids, - std::vector & codepoints, std::string & error) const; + bool extract_fast_tokens(const std::vector & generated_ids, std::vector & fast_ids, + std::string & error) const; - bool decode_fast_tokens(const std::vector> & batch_fast_ids, - FastDecodeResult & result, std::string & error) const; + bool decode_fast_tokens(const std::vector> & batch_fast_ids, FastDecodeResult & result, + std::string & error) const; // Strict low-level API: every input ID must be an action token. Use // decode_generated_tokens for complete Qwen sequences containing text. - bool decode_vlm_action_tokens(const std::vector> & batch_vlm_ids, - FastDecodeResult & result, std::string & error) const; + bool decode_vlm_action_tokens(const std::vector> & batch_vlm_ids, FastDecodeResult & result, + std::string & error) const; // Production entry point for complete Qwen generated_ids. Ordinary text and // control tokens are filtered through the explicit inverse action-token map. bool decode_generated_tokens(const std::vector> & batch_generated_ids, FastDecodeResult & result, std::string & error) const; -private: + private: FastCodec(FastCodecConfig config, std::vector> token_bytes, std::vector fast_to_vlm_id); + bool byte_level_decode(const std::vector & fast_ids, std::vector & codepoints, + std::string & error) const; + FastCodecConfig config_; std::vector> token_bytes_; std::vector fast_to_vlm_id_; diff --git a/src/models/starvla/fast_policy.cpp b/src/models/starvla/fast_policy.cpp index 41e2328..e416fbb 100644 --- a/src/models/starvla/fast_policy.cpp +++ b/src/models/starvla/fast_policy.cpp @@ -2,6 +2,7 @@ #include "ggml.h" #include "gguf.h" +#include "models/starvla/policy_gguf.h" #include #include @@ -17,192 +18,79 @@ namespace robotcpp::starvla { namespace { -constexpr const char * kArchitecture = "starvla-policy"; -constexpr const char * kActionMapTensor = - "starvla.policy.fast.action_token_map"; -constexpr const char * kOffsetsTensor = - "starvla.policy.fast.codec.token_offsets"; -constexpr const char * kTokenBytesTensor = - "starvla.policy.fast.codec.token_bytes"; +constexpr const char * kArchitecture = "starvla-policy"; +constexpr const char * kActionMapTensor = "starvla.policy.fast.action_token_map"; +constexpr const char * kOffsetsTensor = "starvla.policy.fast.codec.token_offsets"; +constexpr const char * kTokenBytesTensor = "starvla.policy.fast.codec.token_bytes"; -int require_key(gguf_context * gguf, const char * key, gguf_type type) { - const int index = gguf_find_key(gguf, key); - if (index < 0) { - throw std::runtime_error(std::string("missing required FAST GGUF metadata: ") + - key); - } - if (gguf_get_kv_type(gguf, index) != type) { - throw std::runtime_error(std::string("invalid FAST GGUF metadata type: ") + - key); - } - return index; -} - -std::string require_string(gguf_context * gguf, const char * key) { - return gguf_get_val_str(gguf, require_key(gguf, key, GGUF_TYPE_STRING)); -} - -int32_t require_i32(gguf_context * gguf, const char * key) { - return gguf_get_val_i32(gguf, require_key(gguf, key, GGUF_TYPE_INT32)); -} - -float require_f32(gguf_context * gguf, const char * key) { - return gguf_get_val_f32(gguf, require_key(gguf, key, GGUF_TYPE_FLOAT32)); -} - -bool require_bool(gguf_context * gguf, const char * key) { - return gguf_get_val_bool(gguf, require_key(gguf, key, GGUF_TYPE_BOOL)); -} - -int require_array(gguf_context * gguf, const char * key, gguf_type type) { - const int index = require_key(gguf, key, GGUF_TYPE_ARRAY); - if (gguf_get_arr_type(gguf, index) != type) { - throw std::runtime_error( - std::string("invalid FAST GGUF array element type: ") + key); - } - return index; -} - -std::vector require_i32_array(gguf_context * gguf, const char * key) { - const int index = require_array(gguf, key, GGUF_TYPE_INT32); - const size_t count = gguf_get_arr_n(gguf, index); - const auto * data = - static_cast(gguf_get_arr_data(gguf, index)); - if (count != 0 && data == nullptr) { - throw std::runtime_error(std::string("missing FAST GGUF array data: ") + - key); - } - return std::vector(data, data + count); -} - -std::vector require_f32_array(gguf_context * gguf, const char * key) { - const int index = require_array(gguf, key, GGUF_TYPE_FLOAT32); - const size_t count = gguf_get_arr_n(gguf, index); - const auto * data = - static_cast(gguf_get_arr_data(gguf, index)); - if (count != 0 && data == nullptr) { - throw std::runtime_error(std::string("missing FAST GGUF array data: ") + - key); - } - return std::vector(data, data + count); -} - -std::vector require_bool_array(gguf_context * gguf, - const char * key) { - const int index = require_array(gguf, key, GGUF_TYPE_BOOL); - const size_t count = gguf_get_arr_n(gguf, index); - const auto * data = - static_cast(gguf_get_arr_data(gguf, index)); - if (count != 0 && data == nullptr) { - throw std::runtime_error(std::string("missing FAST GGUF array data: ") + - key); - } - std::vector values(count); - for (size_t i = 0; i < count; ++i) { - values[i] = data[i] != 0 ? uint8_t{1} : uint8_t{0}; - } - return values; -} - -std::vector require_string_array(gguf_context * gguf, - const char * key) { - const int index = require_array(gguf, key, GGUF_TYPE_STRING); - const size_t count = gguf_get_arr_n(gguf, index); - std::vector values; - values.reserve(count); - for (size_t i = 0; i < count; ++i) { - values.emplace_back(gguf_get_arr_str(gguf, index, i)); - } - return values; -} - -std::string profile_key(int index, const char * suffix) { - return "starvla.normalization.profile." + std::to_string(index) + "." + - suffix; -} +using detail::require_f32; +using detail::require_i32; +using detail::require_i32_array; +using detail::require_string; +using detail::require_string_array; struct FastRuntimeMetadata { FastCodecConfig codec; int token_bytes_count = 0; }; -FastRuntimeMetadata parse_metadata(gguf_context * gguf, - FastPolicyConfig & config) { +FastRuntimeMetadata parse_metadata(gguf_context * gguf, FastPolicyConfig & config) { if (require_string(gguf, "general.architecture") != kArchitecture || - require_i32(gguf, "starvla.schema_version") != 1 || - require_string(gguf, "starvla.framework") != "fast") { + require_i32(gguf, "starvla.schema_version") != 1 || require_string(gguf, "starvla.framework") != "fast") { throw std::runtime_error("GGUF is not a supported StarVLA FAST policy"); } - config.backbone_arch = require_string(gguf, "starvla.backbone.arch"); - config.bundle_uuid = require_string(gguf, "starvla.bundle.uuid"); - config.text_filename = require_string(gguf, "starvla.component.text.filename"); + config.backbone_arch = require_string(gguf, "starvla.backbone.arch"); + config.bundle_uuid = require_string(gguf, "starvla.bundle.uuid"); + config.text_filename = require_string(gguf, "starvla.component.text.filename"); config.mmproj_filename = require_string(gguf, "starvla.component.mmproj.filename"); - if (config.backbone_arch != "qwen2_5_vl" || config.bundle_uuid.empty() || - config.text_filename.empty() || config.mmproj_filename.empty()) { + if (config.backbone_arch != "qwen2_5_vl" || config.bundle_uuid.empty() || config.text_filename.empty() || + config.mmproj_filename.empty()) { throw std::runtime_error("StarVLA FAST bundle metadata is incomplete"); } - config.qwen_hidden_dim = require_i32(gguf, "starvla.qwen.hidden_size"); - config.qwen_input_embedding_dim = - require_i32(gguf, "starvla.qwen.input_embedding_size"); - config.qwen_vocab_size = require_i32(gguf, "starvla.qwen.vocab_size"); - config.qwen_layer_count = require_i32(gguf, "starvla.qwen.layer_count"); - config.cot_template = require_string(gguf, "starvla.prompt.cot_template"); - - config.action_dim = require_i32(gguf, "starvla.action.dimension"); - config.horizon = require_i32(gguf, "starvla.action.horizon"); - config.image_count = require_i32(gguf, "starvla.image.count"); - config.image_names = require_string_array(gguf, "starvla.image.names"); - config.image_processor_min_pixels = - require_i32(gguf, "starvla.image.processor_min_pixels"); - config.image_processor_max_pixels = - require_i32(gguf, "starvla.image.processor_max_pixels"); - config.image_patch_size = require_i32(gguf, "starvla.image.patch_size"); - config.image_spatial_merge_size = - require_i32(gguf, "starvla.image.spatial_merge_size"); - config.image_min_token_count = require_i32(gguf, "starvla.image.min_token_count"); - config.image_max_token_count = require_i32(gguf, "starvla.image.max_token_count"); - - const int max_length = require_i32(gguf, "starvla.fast.generation.max_length"); - config.generation_eos_token_ids = - require_i32_array(gguf, "starvla.fast.generation.eos_token_ids"); - config.generation_top_k = require_i32(gguf, "starvla.fast.generation.top_k"); - config.generation_repetition_penalty = - require_f32(gguf, "starvla.fast.generation.repetition_penalty"); + config.qwen_hidden_dim = require_i32(gguf, "starvla.qwen.hidden_size"); + config.qwen_input_embedding_dim = require_i32(gguf, "starvla.qwen.input_embedding_size"); + config.qwen_vocab_size = require_i32(gguf, "starvla.qwen.vocab_size"); + config.qwen_layer_count = require_i32(gguf, "starvla.qwen.layer_count"); + config.cot_template = require_string(gguf, "starvla.prompt.cot_template"); + + config.action_dim = require_i32(gguf, "starvla.action.dimension"); + config.horizon = require_i32(gguf, "starvla.action.horizon"); + config.image_count = require_i32(gguf, "starvla.image.count"); + config.image_names = require_string_array(gguf, "starvla.image.names"); + config.image_processor_min_pixels = require_i32(gguf, "starvla.image.processor_min_pixels"); + config.image_processor_max_pixels = require_i32(gguf, "starvla.image.processor_max_pixels"); + config.image_patch_size = require_i32(gguf, "starvla.image.patch_size"); + config.image_spatial_merge_size = require_i32(gguf, "starvla.image.spatial_merge_size"); + config.image_min_token_count = require_i32(gguf, "starvla.image.min_token_count"); + config.image_max_token_count = require_i32(gguf, "starvla.image.max_token_count"); + + const int max_length = require_i32(gguf, "starvla.fast.generation.max_length"); + config.generation_eos_token_ids = require_i32_array(gguf, "starvla.fast.generation.eos_token_ids"); + config.generation_top_k = require_i32(gguf, "starvla.fast.generation.top_k"); + config.generation_repetition_penalty = require_f32(gguf, "starvla.fast.generation.repetition_penalty"); FastRuntimeMetadata runtime; - runtime.codec.scale = require_f32(gguf, "starvla.fast.codec.scale"); - runtime.codec.min_token = require_i32(gguf, "starvla.fast.codec.min_token"); - runtime.codec.vocab_size = - static_cast(require_i32(gguf, "starvla.fast.codec.vocab_size")); - runtime.codec.time_horizon = - static_cast(require_i32(gguf, "starvla.fast.codec.time_horizon")); - runtime.codec.action_dim = - static_cast(require_i32(gguf, "starvla.fast.codec.action_dimension")); - const int action_token_count = - require_i32(gguf, "starvla.fast.action_token.count"); - const int offsets_count = - require_i32(gguf, "starvla.fast.codec.token_offsets_count"); - runtime.token_bytes_count = - require_i32(gguf, "starvla.fast.codec.token_bytes_count"); + runtime.codec.scale = require_f32(gguf, "starvla.fast.codec.scale"); + runtime.codec.min_token = require_i32(gguf, "starvla.fast.codec.min_token"); + runtime.codec.vocab_size = static_cast(require_i32(gguf, "starvla.fast.codec.vocab_size")); + runtime.codec.time_horizon = static_cast(require_i32(gguf, "starvla.fast.codec.time_horizon")); + runtime.codec.action_dim = static_cast(require_i32(gguf, "starvla.fast.codec.action_dimension")); + const int action_token_count = require_i32(gguf, "starvla.fast.action_token.count"); + const int offsets_count = require_i32(gguf, "starvla.fast.codec.token_offsets_count"); + runtime.token_bytes_count = require_i32(gguf, "starvla.fast.codec.token_bytes_count"); const bool valid = - config.qwen_hidden_dim > 0 && config.qwen_input_embedding_dim > 0 && - config.qwen_vocab_size > 0 && config.qwen_layer_count > 0 && - !config.cot_template.empty() && config.action_dim > 0 && config.horizon > 0 && - config.image_count > 0 && - config.image_names.size() == static_cast(config.image_count) && + config.qwen_hidden_dim > 0 && config.qwen_input_embedding_dim > 0 && config.qwen_vocab_size > 0 && + config.qwen_layer_count > 0 && !config.cot_template.empty() && config.action_dim > 0 && config.horizon > 0 && + config.image_count > 0 && config.image_names.size() == static_cast(config.image_count) && config.image_processor_min_pixels > 0 && - config.image_processor_max_pixels >= config.image_processor_min_pixels && - config.image_patch_size > 0 && config.image_spatial_merge_size > 0 && - config.image_min_token_count > 0 && - config.image_max_token_count >= config.image_min_token_count && - max_length > 0 && !config.generation_eos_token_ids.empty() && - config.generation_top_k > 0 && - std::isfinite(config.generation_repetition_penalty) && - config.generation_repetition_penalty > 0.0f && - runtime.codec.vocab_size > 0 && - action_token_count == static_cast(runtime.codec.vocab_size) && + config.image_processor_max_pixels >= config.image_processor_min_pixels && config.image_patch_size > 0 && + config.image_spatial_merge_size > 0 && config.image_min_token_count > 0 && + config.image_max_token_count >= config.image_min_token_count && max_length > 0 && + !config.generation_eos_token_ids.empty() && config.generation_top_k > 0 && + std::isfinite(config.generation_repetition_penalty) && config.generation_repetition_penalty > 0.0f && + runtime.codec.vocab_size > 0 && action_token_count == static_cast(runtime.codec.vocab_size) && offsets_count == action_token_count + 1 && runtime.token_bytes_count > 0 && runtime.codec.time_horizon == static_cast(config.horizon) && runtime.codec.action_dim == static_cast(config.action_dim); @@ -211,85 +99,39 @@ FastRuntimeMetadata parse_metadata(gguf_context * gguf, } config.generation_max_length = static_cast(max_length); - NormalizationConfig & normalization = config.normalization; - normalization.clip_actions = require_bool(gguf, "starvla.normalization.clip_actions"); - normalization.binary_threshold = - require_f32(gguf, "starvla.normalization.binary_threshold"); - normalization.binary_comparison = - require_string(gguf, "starvla.normalization.binary_comparison"); - normalization.continuous_dimensions = - require_i32_array(gguf, "starvla.action.continuous_dimensions"); - normalization.binary_dimensions = - require_i32_array(gguf, "starvla.action.binary_dimensions"); - const int profile_count = require_i32(gguf, "starvla.normalization.profile_count"); - const std::vector profile_keys = - require_string_array(gguf, "starvla.normalization.profile_keys"); - if (profile_count <= 0 || profile_keys.size() != static_cast(profile_count)) { - throw std::runtime_error("StarVLA FAST normalization profiles are inconsistent"); - } - normalization.profiles.clear(); - normalization.profiles.reserve(static_cast(profile_count)); - for (int index = 0; index < profile_count; ++index) { - NormalizationProfile profile; - profile.key = require_string(gguf, profile_key(index, "key").c_str()); - profile.action_q01 = - require_f32_array(gguf, profile_key(index, "action_q01").c_str()); - profile.action_q99 = - require_f32_array(gguf, profile_key(index, "action_q99").c_str()); - profile.action_mask = - require_bool_array(gguf, profile_key(index, "action_mask").c_str()); - if (profile.key != profile_keys[static_cast(index)]) { - throw std::runtime_error("StarVLA FAST normalization profile order is inconsistent"); - } - normalization.profiles.push_back(std::move(profile)); - } - std::string normalization_error; - if (!validate_normalization_config(normalization, config.action_dim, - normalization_error)) { - throw std::runtime_error(normalization_error); - } + config.normalization = detail::require_normalization(gguf, config.action_dim); return runtime; } struct RawTensor { ggml_tensor * metadata = nullptr; - int index = -1; + int index = -1; std::vector bytes; }; -RawTensor read_tensor(const std::string & path, gguf_context * gguf, - ggml_context * metadata_context, const char * name, +RawTensor read_tensor(const std::string & path, gguf_context * gguf, ggml_context * metadata_context, const char * name, ggml_type expected_type, int64_t expected_elements) { RawTensor result; result.metadata = ggml_get_tensor(metadata_context, name); - result.index = gguf_find_tensor(gguf, name); - if (result.metadata == nullptr || result.index < 0 || - result.metadata->type != expected_type || - ggml_n_dims(result.metadata) != 1 || - result.metadata->ne[0] != expected_elements || + result.index = gguf_find_tensor(gguf, name); + if (result.metadata == nullptr || result.index < 0 || result.metadata->type != expected_type || + ggml_n_dims(result.metadata) != 1 || result.metadata->ne[0] != expected_elements || ggml_nelements(result.metadata) != expected_elements) { - throw std::runtime_error(std::string("FAST runtime tensor shape/type mismatch: ") + - name); + throw std::runtime_error(std::string("FAST runtime tensor shape/type mismatch: ") + name); } result.bytes.resize(ggml_nbytes(result.metadata)); std::ifstream stream(path, std::ios::binary); if (!stream) { throw std::runtime_error("failed to open FAST policy GGUF tensor data"); } - const size_t offset = - gguf_get_data_offset(gguf) + - gguf_get_tensor_offset(gguf, result.index); + const size_t offset = gguf_get_data_offset(gguf) + gguf_get_tensor_offset(gguf, result.index); stream.seekg(static_cast(offset), std::ios::beg); - if (!stream || - offset > static_cast(std::numeric_limits::max())) { - throw std::runtime_error( - std::string("failed to seek FAST runtime tensor: ") + name); + if (!stream || offset > static_cast(std::numeric_limits::max())) { + throw std::runtime_error(std::string("failed to seek FAST runtime tensor: ") + name); } - stream.read(reinterpret_cast(result.bytes.data()), - static_cast(result.bytes.size())); + stream.read(reinterpret_cast(result.bytes.data()), static_cast(result.bytes.size())); if (!stream) { - throw std::runtime_error( - std::string("failed to read FAST runtime tensor: ") + name); + throw std::runtime_error(std::string("failed to read FAST runtime tensor: ") + name); } return result; } @@ -305,9 +147,7 @@ FastPolicy::FastPolicy(std::unique_ptr impl) : impl_(std::move(impl)) {} FastPolicy::~FastPolicy() = default; -std::unique_ptr FastPolicy::load(const std::string & path, - int verbosity, - std::string & error) { +std::unique_ptr FastPolicy::load(const std::string & path, int verbosity, std::string & error) { error.clear(); if (path.empty()) { error = "StarVLA FAST policy path is required"; @@ -316,8 +156,8 @@ std::unique_ptr FastPolicy::load(const std::string & path, ggml_context * metadata_context = nullptr; gguf_init_params params{}; - params.no_alloc = true; - params.ctx = &metadata_context; + params.no_alloc = true; + params.ctx = &metadata_context; gguf_context * gguf = gguf_init_from_file(path.c_str(), params); if (gguf == nullptr || metadata_context == nullptr) { if (metadata_context != nullptr) { @@ -340,42 +180,32 @@ std::unique_ptr FastPolicy::load(const std::string & path, try { const FastRuntimeMetadata runtime = parse_metadata(gguf, impl->config); - RawTensor action_map = - read_tensor(path, gguf, metadata_context, kActionMapTensor, - GGML_TYPE_I32, static_cast(runtime.codec.vocab_size)); - RawTensor offsets = - read_tensor(path, gguf, metadata_context, kOffsetsTensor, - GGML_TYPE_I32, - static_cast(runtime.codec.vocab_size + 1)); + RawTensor action_map = read_tensor(path, gguf, metadata_context, kActionMapTensor, GGML_TYPE_I32, + static_cast(runtime.codec.vocab_size)); + RawTensor offsets = read_tensor(path, gguf, metadata_context, kOffsetsTensor, GGML_TYPE_I32, + static_cast(runtime.codec.vocab_size + 1)); RawTensor token_bytes = - read_tensor(path, gguf, metadata_context, kTokenBytesTensor, - GGML_TYPE_I8, runtime.token_bytes_count); + read_tensor(path, gguf, metadata_context, kTokenBytesTensor, GGML_TYPE_I8, runtime.token_bytes_count); const uint32_t endian_probe = 1; if (*reinterpret_cast(&endian_probe) != 1) { - throw std::runtime_error( - "FAST runtime currently requires a little-endian host"); + throw std::runtime_error("FAST runtime currently requires a little-endian host"); } std::vector action_ids(runtime.codec.vocab_size); std::vector token_offsets(runtime.codec.vocab_size + 1); - std::memcpy(action_ids.data(), action_map.bytes.data(), - action_map.bytes.size()); - std::memcpy(token_offsets.data(), offsets.bytes.data(), - offsets.bytes.size()); - impl->codec = FastCodec::create_compiled( - runtime.codec, std::move(token_offsets), - std::move(token_bytes.bytes), std::move(action_ids), error); + std::memcpy(action_ids.data(), action_map.bytes.data(), action_map.bytes.size()); + std::memcpy(token_offsets.data(), offsets.bytes.data(), offsets.bytes.size()); + impl->codec = FastCodec::create_compiled(runtime.codec, std::move(token_offsets), std::move(token_bytes.bytes), + std::move(action_ids), error); if (impl->codec == nullptr) { - throw std::runtime_error("failed to construct embedded FAST codec: " + - error); + throw std::runtime_error("failed to construct embedded FAST codec: " + error); } if (verbosity >= 1) { std::fprintf(stderr, "%s: bundle=%s runtime_tensors=3 codec_vocab=%zu " "generation_max_length=%zu profiles=%zu\n", - __func__, impl->config.bundle_uuid.c_str(), - runtime.codec.vocab_size, impl->config.generation_max_length, - impl->config.normalization.profiles.size()); + __func__, impl->config.bundle_uuid.c_str(), runtime.codec.vocab_size, + impl->config.generation_max_length, impl->config.normalization.profiles.size()); } cleanup(); } catch (const std::exception & exception) { @@ -386,35 +216,21 @@ std::unique_ptr FastPolicy::load(const std::string & path, return std::unique_ptr(new FastPolicy(std::move(impl))); } -bool FastPolicy::decode_generated( - const std::vector & full_sequence, - std::vector & action_token_ids, - std::vector & fast_token_ids, - std::vector & normalized_actions, - std::string & error) const { - action_token_ids.clear(); - fast_token_ids.clear(); +bool FastPolicy::decode_generated(const std::vector & full_sequence, std::vector & normalized_actions, + std::string & error) const { normalized_actions.clear(); error.clear(); if (impl_ == nullptr || impl_->codec == nullptr) { error = "StarVLA FAST policy is not initialized"; return false; } - if (!impl_->codec->extract_fast_tokens(full_sequence, fast_token_ids, - error) || - !impl_->codec->map_fast_to_vlm(fast_token_ids, action_token_ids, - error)) { - return false; - } FastDecodeResult decoded; - if (!impl_->codec->decode_fast_tokens({fast_token_ids}, decoded, error)) { + if (!impl_->codec->decode_generated_tokens({full_sequence}, decoded, error)) { return false; } - if (decoded.batch_size != 1 || - decoded.time_horizon != static_cast(impl_->config.horizon) || + if (decoded.batch_size != 1 || decoded.time_horizon != static_cast(impl_->config.horizon) || decoded.action_dim != static_cast(impl_->config.action_dim) || - decoded.actions.size() != - static_cast(impl_->config.horizon * impl_->config.action_dim)) { + decoded.actions.size() != static_cast(impl_->config.horizon * impl_->config.action_dim)) { error = "embedded FAST codec returned an incompatible action tensor"; return false; } @@ -431,17 +247,14 @@ bool FastPolicy::decode_generated( return true; } -bool FastPolicy::unnormalize( - const std::vector & normalized_actions, - const std::string & profile_key, std::vector & actions, - std::string & error) const { +bool FastPolicy::unnormalize(const std::vector & normalized_actions, const std::string & profile_key, + std::vector & actions, std::string & error) const { if (impl_ == nullptr) { actions.clear(); error = "StarVLA FAST policy is not initialized"; return false; } - return denormalize_actions(impl_->config.normalization, profile_key, - normalized_actions, impl_->config.horizon, + return denormalize_actions(impl_->config.normalization, profile_key, normalized_actions, impl_->config.horizon, impl_->config.action_dim, actions, error); } diff --git a/src/models/starvla/fast_policy.h b/src/models/starvla/fast_policy.h index 6a23c11..a92bbfb 100644 --- a/src/models/starvla/fast_policy.h +++ b/src/models/starvla/fast_policy.h @@ -17,27 +17,27 @@ struct FastPolicyConfig { std::string text_filename; std::string mmproj_filename; - int qwen_hidden_dim = 0; + int qwen_hidden_dim = 0; int qwen_input_embedding_dim = 0; - int qwen_vocab_size = 0; - int qwen_layer_count = 0; + int qwen_vocab_size = 0; + int qwen_layer_count = 0; std::string cot_template; int action_dim = 0; - int horizon = 0; + int horizon = 0; int image_count = 0; std::vector image_names; int image_processor_min_pixels = 0; int image_processor_max_pixels = 0; - int image_patch_size = 0; - int image_spatial_merge_size = 0; - int image_min_token_count = 0; - int image_max_token_count = 0; + int image_patch_size = 0; + int image_spatial_merge_size = 0; + int image_min_token_count = 0; + int image_max_token_count = 0; size_t generation_max_length = 0; std::vector generation_eos_token_ids; - int generation_top_k = 0; + int generation_top_k = 0; float generation_repetition_penalty = 0.0f; NormalizationConfig normalization; @@ -47,23 +47,16 @@ class FastPolicy { public: ~FastPolicy(); - FastPolicy(const FastPolicy &) = delete; + FastPolicy(const FastPolicy &) = delete; FastPolicy & operator=(const FastPolicy &) = delete; - static std::unique_ptr load(const std::string & path, - int verbosity, - std::string & error); + static std::unique_ptr load(const std::string & path, int verbosity, std::string & error); - bool decode_generated(const std::vector & full_sequence, - std::vector & action_token_ids, - std::vector & fast_token_ids, - std::vector & normalized_actions, + bool decode_generated(const std::vector & full_sequence, std::vector & normalized_actions, std::string & error) const; - bool unnormalize(const std::vector & normalized_actions, - const std::string & profile_key, - std::vector & actions, - std::string & error) const; + bool unnormalize(const std::vector & normalized_actions, const std::string & profile_key, + std::vector & actions, std::string & error) const; const FastPolicyConfig & config() const; const char * backend_name() const; diff --git a/src/models/starvla/groot_policy.cpp b/src/models/starvla/groot_policy.cpp index d9c1296..a91a150 100644 --- a/src/models/starvla/groot_policy.cpp +++ b/src/models/starvla/groot_policy.cpp @@ -5,11 +5,11 @@ #include "gguf.h" #include "models/ggml_backend.h" #include "models/gguf_loader.h" +#include "models/starvla/policy_gguf.h" #include #include #include -#include #include #include #include @@ -21,157 +21,55 @@ namespace robotcpp::starvla { namespace { constexpr size_t kGraphSize = 16384; -constexpr int kKQMaskPad = 32; +constexpr int kKQMaskPad = 32; struct GR00TBlockWeights { - ggml_tensor * ada_norm_weight = nullptr; - ggml_tensor * ada_norm_bias = nullptr; - ggml_tensor * query_weight = nullptr; - ggml_tensor * query_bias = nullptr; - ggml_tensor * key_weight = nullptr; - ggml_tensor * key_bias = nullptr; - ggml_tensor * value_weight = nullptr; - ggml_tensor * value_bias = nullptr; - ggml_tensor * attention_output_weight = nullptr; - ggml_tensor * attention_output_bias = nullptr; - ggml_tensor * feed_forward_input_weight = nullptr; - ggml_tensor * feed_forward_input_bias = nullptr; + ggml_tensor * ada_norm_weight = nullptr; + ggml_tensor * ada_norm_bias = nullptr; + ggml_tensor * query_weight = nullptr; + ggml_tensor * query_bias = nullptr; + ggml_tensor * key_weight = nullptr; + ggml_tensor * key_bias = nullptr; + ggml_tensor * value_weight = nullptr; + ggml_tensor * value_bias = nullptr; + ggml_tensor * attention_output_weight = nullptr; + ggml_tensor * attention_output_bias = nullptr; + ggml_tensor * feed_forward_input_weight = nullptr; + ggml_tensor * feed_forward_input_bias = nullptr; ggml_tensor * feed_forward_output_weight = nullptr; - ggml_tensor * feed_forward_output_bias = nullptr; + ggml_tensor * feed_forward_output_bias = nullptr; }; struct GR00TWeights { - ggml_tensor * timestep_input_weight = nullptr; - ggml_tensor * timestep_input_bias = nullptr; + ggml_tensor * timestep_input_weight = nullptr; + ggml_tensor * timestep_input_bias = nullptr; ggml_tensor * timestep_output_weight = nullptr; - ggml_tensor * timestep_output_bias = nullptr; + ggml_tensor * timestep_output_bias = nullptr; std::vector blocks; ggml_tensor * output_modulation_weight = nullptr; - ggml_tensor * output_modulation_bias = nullptr; + ggml_tensor * output_modulation_bias = nullptr; ggml_tensor * output_projection_weight = nullptr; - ggml_tensor * output_projection_bias = nullptr; - ggml_tensor * action_input_weight = nullptr; - ggml_tensor * action_input_bias = nullptr; - ggml_tensor * action_time_mix_weight = nullptr; - ggml_tensor * action_time_mix_bias = nullptr; - ggml_tensor * action_output_weight = nullptr; - ggml_tensor * action_output_bias = nullptr; - ggml_tensor * velocity_input_weight = nullptr; - ggml_tensor * velocity_input_bias = nullptr; - ggml_tensor * velocity_output_weight = nullptr; - ggml_tensor * velocity_output_bias = nullptr; - ggml_tensor * future_tokens = nullptr; - ggml_tensor * action_position = nullptr; + ggml_tensor * output_projection_bias = nullptr; + ggml_tensor * action_input_weight = nullptr; + ggml_tensor * action_input_bias = nullptr; + ggml_tensor * action_time_mix_weight = nullptr; + ggml_tensor * action_time_mix_bias = nullptr; + ggml_tensor * action_output_weight = nullptr; + ggml_tensor * action_output_bias = nullptr; + ggml_tensor * velocity_input_weight = nullptr; + ggml_tensor * velocity_input_bias = nullptr; + ggml_tensor * velocity_output_weight = nullptr; + ggml_tensor * velocity_output_bias = nullptr; + ggml_tensor * future_tokens = nullptr; + ggml_tensor * action_position = nullptr; }; -int require_key(gguf_context * gguf, const char * key, gguf_type type) { - const int index = gguf_find_key(gguf, key); - if (index < 0) { - throw std::runtime_error(std::string("missing required StarVLA GGUF metadata: ") + key); - } - if (gguf_get_kv_type(gguf, index) != type) { - throw std::runtime_error(std::string("invalid StarVLA GGUF metadata type: ") + key); - } - return index; -} - -std::string require_string(gguf_context * gguf, const char * key) { - return gguf_get_val_str(gguf, require_key(gguf, key, GGUF_TYPE_STRING)); -} - -int require_i32(gguf_context * gguf, const char * key) { - return gguf_get_val_i32(gguf, require_key(gguf, key, GGUF_TYPE_INT32)); -} - -float require_f32(gguf_context * gguf, const char * key) { - return gguf_get_val_f32(gguf, require_key(gguf, key, GGUF_TYPE_FLOAT32)); -} - -bool require_bool(gguf_context * gguf, const char * key) { - return gguf_get_val_bool(gguf, require_key(gguf, key, GGUF_TYPE_BOOL)); -} - -int require_array(gguf_context * gguf, const char * key, gguf_type element_type) { - const int index = require_key(gguf, key, GGUF_TYPE_ARRAY); - if (gguf_get_arr_type(gguf, index) != element_type) { - throw std::runtime_error(std::string("invalid StarVLA GGUF array element type: ") + key); - } - return index; -} - -std::vector require_string_array(gguf_context * gguf, const char * key) { - const int index = require_array(gguf, key, GGUF_TYPE_STRING); - const size_t count = gguf_get_arr_n(gguf, index); - std::vector result; - result.reserve(count); - for (size_t i = 0; i < count; ++i) { - result.emplace_back(gguf_get_arr_str(gguf, index, i)); - } - return result; -} - -std::vector require_i32_array(gguf_context * gguf, const char * key) { - const int index = require_array(gguf, key, GGUF_TYPE_INT32); - const size_t count = gguf_get_arr_n(gguf, index); - const auto * data = static_cast(gguf_get_arr_data(gguf, index)); - if (data == nullptr && count != 0) { - throw std::runtime_error(std::string("missing StarVLA GGUF array data: ") + key); - } - return count == 0 ? std::vector() : std::vector(data, data + count); -} - -std::vector require_f32_array(gguf_context * gguf, const char * key) { - const int index = require_array(gguf, key, GGUF_TYPE_FLOAT32); - const size_t count = gguf_get_arr_n(gguf, index); - const auto * data = static_cast(gguf_get_arr_data(gguf, index)); - if (data == nullptr && count != 0) { - throw std::runtime_error(std::string("missing StarVLA GGUF array data: ") + key); - } - return count == 0 ? std::vector() : std::vector(data, data + count); -} - -std::vector require_bool_array(gguf_context * gguf, const char * key) { - const int index = require_array(gguf, key, GGUF_TYPE_BOOL); - const size_t count = gguf_get_arr_n(gguf, index); - const auto * data = static_cast(gguf_get_arr_data(gguf, index)); - if (data == nullptr && count != 0) { - throw std::runtime_error(std::string("missing StarVLA GGUF array data: ") + key); - } - std::vector result(count); - for (size_t i = 0; i < count; ++i) { - result[i] = data[i] != 0 ? 1 : 0; - } - return result; -} - -std::string profile_key(int profile_index, const char * suffix) { - return "starvla.normalization.profile." + std::to_string(profile_index) + "." + suffix; -} - -bool has_shape(const ggml_tensor * tensor, std::initializer_list expected) { - if (tensor == nullptr || static_cast(ggml_n_dims(tensor)) != expected.size()) { - return false; - } - size_t dimension = 0; - for (const int64_t value : expected) { - if (tensor->ne[dimension++] != value) { - return false; - } - } - return true; -} - -const char * mode_name(backend_mode mode) { - switch (mode) { - case backend_mode::cpu: - return "cpu"; - case backend_mode::cuda: - return "cuda"; - case backend_mode::metal: - return "metal"; - } - return "unknown"; -} +using detail::has_shape; +using detail::require_f32; +using detail::require_i32; +using detail::require_i32_array; +using detail::require_string; +using detail::require_string_array; class GR00TGGUFLoader final : public gguf_loader { public: @@ -182,130 +80,84 @@ class GR00TGGUFLoader final : public gguf_loader { if (require_string(gguf, "general.architecture") != "starvla-policy") { throw std::runtime_error("StarVLA GR00T policy has incompatible general.architecture"); } - if (require_i32(gguf, "starvla.schema_version") != 1 || - require_string(gguf, "starvla.framework") != "groot") { + if (require_i32(gguf, "starvla.schema_version") != 1 || require_string(gguf, "starvla.framework") != "groot") { throw std::runtime_error("StarVLA policy GGUF is not a supported Qwen GR00T schema"); } config_.backbone_arch = require_string(gguf, "starvla.backbone.arch"); - if (config_.backbone_arch != "qwen3_vl" && - config_.backbone_arch != "qwen2_5_vl") { - throw std::runtime_error( - "StarVLA GR00T policy has an unsupported Qwen backbone"); + if (config_.backbone_arch != "qwen3_vl" && config_.backbone_arch != "qwen2_5_vl") { + throw std::runtime_error("StarVLA GR00T policy has an unsupported Qwen backbone"); } config_.bundle_uuid = require_string(gguf, "starvla.bundle.uuid"); if (config_.bundle_uuid.empty()) { throw std::runtime_error("StarVLA GR00T bundle UUID is missing"); } - config_.text_filename = require_string(gguf, "starvla.component.text.filename"); + config_.text_filename = require_string(gguf, "starvla.component.text.filename"); config_.mmproj_filename = require_string(gguf, "starvla.component.mmproj.filename"); if (config_.text_filename.empty() || config_.mmproj_filename.empty()) { throw std::runtime_error("StarVLA GR00T component filenames must be non-empty"); } - config_.qwen_hidden_dim = require_i32(gguf, "starvla.qwen.hidden_size"); - config_.qwen_input_embedding_dim = - require_i32(gguf, "starvla.qwen.input_embedding_size"); - config_.qwen_vocab_size = require_i32(gguf, "starvla.qwen.vocab_size"); - config_.cot_template = require_string(gguf, "starvla.prompt.cot_template"); - const bool qwen25 = config_.backbone_arch == "qwen2_5_vl"; + config_.qwen_hidden_dim = require_i32(gguf, "starvla.qwen.hidden_size"); + config_.qwen_input_embedding_dim = require_i32(gguf, "starvla.qwen.input_embedding_size"); + config_.qwen_vocab_size = require_i32(gguf, "starvla.qwen.vocab_size"); + config_.cot_template = require_string(gguf, "starvla.prompt.cot_template"); + const bool qwen25 = config_.backbone_arch == "qwen2_5_vl"; if (config_.cot_template.empty()) { throw std::runtime_error("StarVLA GR00T prompt template is missing"); } - config_.image_count = require_i32(gguf, "starvla.image.count"); - config_.image_names = require_string_array(gguf, "starvla.image.names"); - config_.image_processor_min_pixels = - require_i32(gguf, "starvla.image.processor_min_pixels"); - config_.image_processor_max_pixels = - require_i32(gguf, "starvla.image.processor_max_pixels"); - config_.image_patch_size = require_i32(gguf, "starvla.image.patch_size"); - config_.image_spatial_merge_size = - require_i32(gguf, "starvla.image.spatial_merge_size"); - config_.image_min_token_count = - require_i32(gguf, "starvla.image.min_token_count"); - config_.image_max_token_count = - require_i32(gguf, "starvla.image.max_token_count"); - config_.dit_width = require_i32(gguf, "starvla.groot.dit_width"); - config_.block_count = require_i32(gguf, "starvla.groot.block_count"); - config_.attention_head_count = require_i32(gguf, "starvla.groot.attention_head_count"); - config_.attention_head_dim = require_i32(gguf, "starvla.groot.attention_head_dim"); - config_.cross_attention_dim = require_i32(gguf, "starvla.groot.cross_attention_dim"); - config_.feed_forward_dim = require_i32(gguf, "starvla.groot.feed_forward_dim"); - config_.output_dim = require_i32(gguf, "starvla.groot.output_dimension"); - config_.mlp_hidden_dim = require_i32(gguf, "starvla.groot.mlp_hidden_dimension"); - config_.future_token_count = require_i32(gguf, "starvla.groot.future_token_count"); - config_.action_position_count = require_i32(gguf, "starvla.groot.action_position_count"); - config_.no_state_sequence_length = require_i32(gguf, "starvla.groot.no_state_sequence_length"); - config_.timestep_projection_dim = require_i32(gguf, "starvla.groot.timestep_projection_dim"); - config_.ada_norm_epsilon = require_f32(gguf, "starvla.groot.ada_norm_epsilon"); - config_.output_norm_epsilon = require_f32(gguf, "starvla.groot.output_norm_epsilon"); - config_.euler_dt = require_f32(gguf, "starvla.groot.euler_dt"); - config_.timestep_ids = require_i32_array(gguf, "starvla.groot.timestep_ids"); - config_.action_dim = require_i32(gguf, "starvla.action.dimension"); - config_.horizon = require_i32(gguf, "starvla.action.horizon"); + config_.image_count = require_i32(gguf, "starvla.image.count"); + config_.image_names = require_string_array(gguf, "starvla.image.names"); + config_.image_processor_min_pixels = require_i32(gguf, "starvla.image.processor_min_pixels"); + config_.image_processor_max_pixels = require_i32(gguf, "starvla.image.processor_max_pixels"); + config_.image_patch_size = require_i32(gguf, "starvla.image.patch_size"); + config_.image_spatial_merge_size = require_i32(gguf, "starvla.image.spatial_merge_size"); + config_.image_min_token_count = require_i32(gguf, "starvla.image.min_token_count"); + config_.image_max_token_count = require_i32(gguf, "starvla.image.max_token_count"); + config_.dit_width = require_i32(gguf, "starvla.groot.dit_width"); + config_.block_count = require_i32(gguf, "starvla.groot.block_count"); + config_.attention_head_count = require_i32(gguf, "starvla.groot.attention_head_count"); + config_.attention_head_dim = require_i32(gguf, "starvla.groot.attention_head_dim"); + config_.cross_attention_dim = require_i32(gguf, "starvla.groot.cross_attention_dim"); + config_.feed_forward_dim = require_i32(gguf, "starvla.groot.feed_forward_dim"); + config_.output_dim = require_i32(gguf, "starvla.groot.output_dimension"); + config_.mlp_hidden_dim = require_i32(gguf, "starvla.groot.mlp_hidden_dimension"); + config_.future_token_count = require_i32(gguf, "starvla.groot.future_token_count"); + config_.action_position_count = require_i32(gguf, "starvla.groot.action_position_count"); + config_.no_state_sequence_length = require_i32(gguf, "starvla.groot.no_state_sequence_length"); + config_.timestep_projection_dim = require_i32(gguf, "starvla.groot.timestep_projection_dim"); + config_.ada_norm_epsilon = require_f32(gguf, "starvla.groot.ada_norm_epsilon"); + config_.output_norm_epsilon = require_f32(gguf, "starvla.groot.output_norm_epsilon"); + config_.euler_dt = require_f32(gguf, "starvla.groot.euler_dt"); + config_.timestep_ids = require_i32_array(gguf, "starvla.groot.timestep_ids"); + config_.action_dim = require_i32(gguf, "starvla.action.dimension"); + config_.horizon = require_i32(gguf, "starvla.action.horizon"); const int64_t expected_input_embedding_dim = - qwen25 ? static_cast(config_.qwen_hidden_dim) - : 4LL * config_.qwen_hidden_dim; + qwen25 ? static_cast(config_.qwen_hidden_dim) : 4LL * config_.qwen_hidden_dim; const bool dimensions_valid = config_.qwen_hidden_dim > 0 && config_.qwen_vocab_size > 0 && config_.dit_width > 0 && - config_.qwen_input_embedding_dim == expected_input_embedding_dim && - config_.dit_width % 2 == 0 && + config_.qwen_input_embedding_dim == expected_input_embedding_dim && config_.dit_width % 2 == 0 && config_.block_count > 0 && config_.block_count % 2 == 0 && config_.attention_head_count > 0 && config_.attention_head_dim > 0 && config_.attention_head_count * config_.attention_head_dim == config_.dit_width && config_.cross_attention_dim == config_.qwen_hidden_dim && config_.feed_forward_dim > 0 && - config_.output_dim > 0 && config_.mlp_hidden_dim > 0 && - config_.action_dim > 0 && config_.horizon > 0 && config_.future_token_count > 0 && - config_.action_position_count >= config_.horizon && + config_.output_dim > 0 && config_.mlp_hidden_dim > 0 && config_.action_dim > 0 && config_.horizon > 0 && + config_.future_token_count > 0 && config_.action_position_count >= config_.horizon && config_.no_state_sequence_length == config_.future_token_count + config_.horizon && config_.timestep_projection_dim >= 4 && config_.timestep_projection_dim % 2 == 0 && std::isfinite(config_.ada_norm_epsilon) && config_.ada_norm_epsilon > 0.0f && std::isfinite(config_.output_norm_epsilon) && config_.output_norm_epsilon > 0.0f && - std::isfinite(config_.euler_dt) && config_.euler_dt > 0.0f && - config_.timestep_ids.size() == 4 && - config_.image_count > 0 && - config_.image_names.size() == static_cast(config_.image_count) && + std::isfinite(config_.euler_dt) && config_.euler_dt > 0.0f && config_.timestep_ids.size() == 4 && + config_.image_count > 0 && config_.image_names.size() == static_cast(config_.image_count) && config_.image_processor_min_pixels > 0 && - config_.image_processor_max_pixels >= config_.image_processor_min_pixels && - config_.image_patch_size > 0 && config_.image_spatial_merge_size > 0 && - config_.image_min_token_count > 0 && + config_.image_processor_max_pixels >= config_.image_processor_min_pixels && config_.image_patch_size > 0 && + config_.image_spatial_merge_size > 0 && config_.image_min_token_count > 0 && config_.image_max_token_count >= config_.image_min_token_count; if (!dimensions_valid) { throw std::runtime_error("StarVLA GR00T policy metadata has incompatible dimensions"); } - NormalizationConfig & normalization = config_.normalization; - normalization.clip_actions = require_bool(gguf, "starvla.normalization.clip_actions"); - normalization.binary_threshold = require_f32(gguf, "starvla.normalization.binary_threshold"); - normalization.binary_comparison = require_string(gguf, "starvla.normalization.binary_comparison"); - normalization.continuous_dimensions = - require_i32_array(gguf, "starvla.action.continuous_dimensions"); - normalization.binary_dimensions = require_i32_array(gguf, "starvla.action.binary_dimensions"); - const int profile_count = require_i32(gguf, "starvla.normalization.profile_count"); - const std::vector keys = - require_string_array(gguf, "starvla.normalization.profile_keys"); - if (profile_count <= 0 || keys.size() != static_cast(profile_count)) { - throw std::runtime_error("StarVLA GR00T normalization profile count is inconsistent"); - } - normalization.profiles.clear(); - for (int profile_index = 0; profile_index < profile_count; ++profile_index) { - NormalizationProfile profile; - profile.key = require_string(gguf, profile_key(profile_index, "key").c_str()); - profile.action_q01 = - require_f32_array(gguf, profile_key(profile_index, "action_q01").c_str()); - profile.action_q99 = - require_f32_array(gguf, profile_key(profile_index, "action_q99").c_str()); - profile.action_mask = - require_bool_array(gguf, profile_key(profile_index, "action_mask").c_str()); - if (profile.key != keys[static_cast(profile_index)]) { - throw std::runtime_error("StarVLA GR00T normalization profile order is inconsistent"); - } - normalization.profiles.push_back(std::move(profile)); - } - std::string normalization_error; - if (!validate_normalization_config(normalization, config_.action_dim, normalization_error)) { - throw std::runtime_error(normalization_error); - } + config_.normalization = detail::require_normalization(gguf, config_.action_dim); return true; } @@ -380,14 +232,11 @@ class GR00TGGUFLoader final : public gguf_loader { } for (int block = 0; block < config_.block_count; ++block) { const GR00TBlockWeights & current = weights_.blocks[static_cast(block)]; - const int kv_input_dim = block % 2 == 0 ? config_.cross_attention_dim : width; + const int kv_input_dim = block % 2 == 0 ? config_.cross_attention_dim : width; if (!has_shape(current.ada_norm_weight, {width, 2 * width}) || - !has_shape(current.ada_norm_bias, {2 * width}) || - !has_shape(current.query_weight, {width, width}) || - !has_shape(current.query_bias, {width}) || - !has_shape(current.key_weight, {kv_input_dim, width}) || - !has_shape(current.key_bias, {width}) || - !has_shape(current.value_weight, {kv_input_dim, width}) || + !has_shape(current.ada_norm_bias, {2 * width}) || !has_shape(current.query_weight, {width, width}) || + !has_shape(current.query_bias, {width}) || !has_shape(current.key_weight, {kv_input_dim, width}) || + !has_shape(current.key_bias, {width}) || !has_shape(current.value_weight, {kv_input_dim, width}) || !has_shape(current.value_bias, {width}) || !has_shape(current.attention_output_weight, {width, width}) || !has_shape(current.attention_output_bias, {width}) || @@ -407,36 +256,36 @@ class GR00TGGUFLoader final : public gguf_loader { }; std::vector timestep_projection_table(const GR00TPolicyConfig & config) { - const int dim = config.timestep_projection_dim; - const int half = dim / 2; + const int dim = config.timestep_projection_dim; + const int half = dim / 2; const float denominator = static_cast(half - 1); std::vector table(static_cast(dim) * 4, 0.0f); for (int step = 0; step < 4; ++step) { const float timestep = static_cast(config.timestep_ids[static_cast(step)]); - float * row = table.data() + static_cast(step) * dim; + float * row = table.data() + static_cast(step) * dim; for (int index = 0; index < half; ++index) { const float frequency = std::exp(-std::log(10000.0f) * static_cast(index) / denominator); - const float angle = timestep * frequency; - row[index] = std::cos(angle); - row[index + half] = std::sin(angle); + const float angle = timestep * frequency; + row[index] = std::cos(angle); + row[index + half] = std::sin(angle); } } return table; } std::vector action_time_table(const GR00TPolicyConfig & config) { - const int dim = config.dit_width; - const int half = dim / 2; + const int dim = config.dit_width; + const int half = dim / 2; const float denominator = static_cast(half); std::vector table(static_cast(dim) * 4, 0.0f); for (int step = 0; step < 4; ++step) { const float timestep = static_cast(config.timestep_ids[static_cast(step)]); - float * row = table.data() + static_cast(step) * dim; + float * row = table.data() + static_cast(step) * dim; for (int index = 0; index < half; ++index) { const float frequency = std::exp(-std::log(10000.0f) * static_cast(index) / denominator); - const float angle = timestep * frequency; - row[index] = std::sin(angle); - row[index + half] = std::cos(angle); + const float angle = timestep * frequency; + row[index] = std::sin(angle); + row[index + half] = std::cos(angle); } } return table; @@ -453,19 +302,19 @@ struct GR00TPolicy::Impl { ggml_backend_sched_t scheduler = nullptr; backend_buft_policy buft_policy; backend_mode mode = backend_mode::cpu; - int n_threads = 0; - int verbosity = 0; - - ggml_context * graph_context = nullptr; - ggml_cgraph * graph = nullptr; - ggml_tensor * hidden_input = nullptr; - ggml_tensor * cross_mask_input = nullptr; - ggml_tensor * noise_input = nullptr; + int n_threads = 0; + int verbosity = 0; + + ggml_context * graph_context = nullptr; + ggml_cgraph * graph = nullptr; + ggml_tensor * hidden_input = nullptr; + ggml_tensor * cross_mask_input = nullptr; + ggml_tensor * noise_input = nullptr; ggml_tensor * timestep_projection_input = nullptr; - ggml_tensor * action_time_input = nullptr; - ggml_tensor * scalar_one_input = nullptr; - ggml_tensor * output = nullptr; - size_t conditioning_token_count = 0; + ggml_tensor * action_time_input = nullptr; + ggml_tensor * scalar_one_input = nullptr; + ggml_tensor * output = nullptr; + size_t conditioning_token_count = 0; std::vector timestep_table; std::vector action_table; @@ -506,15 +355,15 @@ struct GR00TPolicy::Impl { ggml_free(graph_context); graph_context = nullptr; } - graph = nullptr; - hidden_input = nullptr; - cross_mask_input = nullptr; - noise_input = nullptr; + graph = nullptr; + hidden_input = nullptr; + cross_mask_input = nullptr; + noise_input = nullptr; timestep_projection_input = nullptr; - action_time_input = nullptr; - scalar_one_input = nullptr; - output = nullptr; - conditioning_token_count = 0; + action_time_input = nullptr; + scalar_one_input = nullptr; + output = nullptr; + conditioning_token_count = 0; } void build_graph(size_t token_count) { @@ -524,29 +373,28 @@ struct GR00TPolicy::Impl { } ggml_init_params params{}; - params.mem_size = kGraphSize * ggml_tensor_overhead() + ggml_graph_overhead_custom(kGraphSize, false); + params.mem_size = kGraphSize * ggml_tensor_overhead() + ggml_graph_overhead_custom(kGraphSize, false); params.mem_buffer = nullptr; - params.no_alloc = true; - graph_context = ggml_init(params); + params.no_alloc = true; + graph_context = ggml_init(params); if (graph_context == nullptr) { throw std::runtime_error("failed to initialize StarVLA GR00T graph context"); } - const int width = config.dit_width; - const int heads = config.attention_head_count; - const int head_dim = config.attention_head_dim; + const int width = config.dit_width; + const int heads = config.attention_head_count; + const int head_dim = config.attention_head_dim; const int sequence_length = config.no_state_sequence_length; - const int mask_queries = GGML_PAD(sequence_length, kKQMaskPad); - - hidden_input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, config.qwen_hidden_dim, - static_cast(token_count)); - cross_mask_input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, - static_cast(token_count), mask_queries); - noise_input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, config.action_dim, config.horizon); - timestep_projection_input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, - config.timestep_projection_dim, 4); - action_time_input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, width, 4); - scalar_one_input = ggml_new_tensor_1d(graph_context, GGML_TYPE_F32, 1); + const int mask_queries = GGML_PAD(sequence_length, kKQMaskPad); + + hidden_input = + ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, config.qwen_hidden_dim, static_cast(token_count)); + cross_mask_input = + ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, static_cast(token_count), mask_queries); + noise_input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, config.action_dim, config.horizon); + timestep_projection_input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, config.timestep_projection_dim, 4); + action_time_input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, width, 4); + scalar_one_input = ggml_new_tensor_1d(graph_context, GGML_TYPE_F32, 1); if (hidden_input == nullptr || cross_mask_input == nullptr || noise_input == nullptr || timestep_projection_input == nullptr || action_time_input == nullptr || scalar_one_input == nullptr) { throw std::runtime_error("failed to create StarVLA GR00T graph inputs"); @@ -573,96 +421,93 @@ struct GR00TPolicy::Impl { return ggml_add(graph_context, projected, f32(bias)); }; auto ada_norm = [&](ggml_tensor * value, ggml_tensor * temb, const GR00TBlockWeights & block) { - ggml_tensor * modulation = linear(ggml_silu(graph_context, temb), block.ada_norm_weight, - block.ada_norm_bias); + ggml_tensor * modulation = + linear(ggml_silu(graph_context, temb), block.ada_norm_weight, block.ada_norm_bias); ggml_tensor * scale = ggml_view_1d(graph_context, modulation, width, 0); - ggml_tensor * shift = ggml_view_1d(graph_context, modulation, width, - static_cast(width) * sizeof(float)); - ggml_tensor * normalized = ggml_norm(graph_context, value, config.ada_norm_epsilon); + ggml_tensor * shift = + ggml_view_1d(graph_context, modulation, width, static_cast(width) * sizeof(float)); + ggml_tensor * normalized = ggml_norm(graph_context, value, config.ada_norm_epsilon); ggml_tensor * one_plus_scale = ggml_add(graph_context, scale, scalar_one_input); return ggml_add(graph_context, ggml_mul(graph_context, normalized, one_plus_scale), shift); }; - auto attention = [&](ggml_tensor * query_source, ggml_tensor * key_value_source, - ggml_tensor * mask, const GR00TBlockWeights & block) { - const int64_t query_count = query_source->ne[1]; + auto attention = [&](ggml_tensor * query_source, ggml_tensor * key_value_source, ggml_tensor * mask, + const GR00TBlockWeights & block) { + const int64_t query_count = query_source->ne[1]; const int64_t key_value_count = key_value_source->ne[1]; - ggml_tensor * query = linear(query_source, block.query_weight, block.query_bias); - ggml_tensor * key = linear(key_value_source, block.key_weight, block.key_bias); - ggml_tensor * value = linear(key_value_source, block.value_weight, block.value_bias); - query = ggml_reshape_3d(graph_context, query, head_dim, heads, query_count); - key = ggml_reshape_3d(graph_context, key, head_dim, heads, key_value_count); - value = ggml_reshape_3d(graph_context, value, head_dim, heads, key_value_count); - query = ggml_permute(graph_context, query, 0, 2, 1, 3); - key = ggml_permute(graph_context, key, 0, 2, 1, 3); - value = ggml_cont(graph_context, ggml_permute(graph_context, value, 1, 2, 0, 3)); - ggml_tensor * scores = ggml_mul_mat(graph_context, key, query); + ggml_tensor * query = linear(query_source, block.query_weight, block.query_bias); + ggml_tensor * key = linear(key_value_source, block.key_weight, block.key_bias); + ggml_tensor * value = linear(key_value_source, block.value_weight, block.value_bias); + query = ggml_reshape_3d(graph_context, query, head_dim, heads, query_count); + key = ggml_reshape_3d(graph_context, key, head_dim, heads, key_value_count); + value = ggml_reshape_3d(graph_context, value, head_dim, heads, key_value_count); + query = ggml_permute(graph_context, query, 0, 2, 1, 3); + key = ggml_permute(graph_context, key, 0, 2, 1, 3); + value = ggml_cont(graph_context, ggml_permute(graph_context, value, 1, 2, 0, 3)); + ggml_tensor * scores = ggml_mul_mat(graph_context, key, query); ggml_mul_mat_set_prec(scores, GGML_PREC_F32); - scores = ggml_soft_max_ext(graph_context, scores, mask, - 1.0f / std::sqrt(static_cast(head_dim)), 0.0f); + scores = + ggml_soft_max_ext(graph_context, scores, mask, 1.0f / std::sqrt(static_cast(head_dim)), 0.0f); ggml_tensor * attended = ggml_mul_mat(graph_context, value, scores); ggml_mul_mat_set_prec(attended, GGML_PREC_F32); attended = ggml_permute(graph_context, attended, 0, 2, 1, 3); attended = ggml_cont_2d(graph_context, attended, width, query_count); return linear(attended, block.attention_output_weight, block.attention_output_bias); }; - ggml_tensor * future = f32(weights.future_tokens); + ggml_tensor * future = f32(weights.future_tokens); ggml_tensor * position_view = ggml_view_2d(graph_context, weights.action_position, width, config.horizon, weights.action_position->nb[1], 0); - ggml_tensor * position = f32(position_view); - ggml_tensor * actions = noise_input; + ggml_tensor * position = f32(position_view); + ggml_tensor * actions = noise_input; for (int step = 0; step < 4; ++step) { - ggml_tensor * timestep_projection = ggml_view_1d( - graph_context, timestep_projection_input, config.timestep_projection_dim, - static_cast(step) * config.timestep_projection_dim * sizeof(float)); - ggml_tensor * temb = linear(timestep_projection, weights.timestep_input_weight, - weights.timestep_input_bias); + ggml_tensor * timestep_projection = + ggml_view_1d(graph_context, timestep_projection_input, config.timestep_projection_dim, + static_cast(step) * config.timestep_projection_dim * sizeof(float)); + ggml_tensor * temb = + linear(timestep_projection, weights.timestep_input_weight, weights.timestep_input_bias); temb = ggml_silu(graph_context, temb); temb = linear(temb, weights.timestep_output_weight, weights.timestep_output_bias); ggml_tensor * action_features = linear(actions, weights.action_input_weight, weights.action_input_bias); - ggml_tensor * action_time = ggml_view_1d( - graph_context, action_time_input, width, - static_cast(step) * width * sizeof(float)); - action_time = ggml_repeat(graph_context, action_time, action_features); - action_features = ggml_concat(graph_context, action_features, action_time, 0); - action_features = linear(action_features, weights.action_time_mix_weight, - weights.action_time_mix_bias); + ggml_tensor * action_time = ggml_view_1d(graph_context, action_time_input, width, + static_cast(step) * width * sizeof(float)); + action_time = ggml_repeat(graph_context, action_time, action_features); + action_features = ggml_concat(graph_context, action_features, action_time, 0); + action_features = linear(action_features, weights.action_time_mix_weight, weights.action_time_mix_bias); action_features = ggml_silu(graph_context, action_features); action_features = linear(action_features, weights.action_output_weight, weights.action_output_bias); action_features = ggml_add(graph_context, action_features, position); ggml_tensor * hidden = ggml_concat(graph_context, future, action_features, 1); for (int block_index = 0; block_index < config.block_count; ++block_index) { const GR00TBlockWeights & block = weights.blocks[static_cast(block_index)]; - ggml_tensor * normalized = ada_norm(hidden, temb, block); - ggml_tensor * attended = block_index % 2 == 0 - ? attention(normalized, hidden_input, cross_mask_input, block) - : attention(normalized, normalized, nullptr, block); - hidden = ggml_add(graph_context, hidden, attended); - ggml_tensor * ff = ggml_norm(graph_context, hidden, config.ada_norm_epsilon); - ff = linear(ff, block.feed_forward_input_weight, block.feed_forward_input_bias); - ff = ggml_gelu(graph_context, ff); - ff = linear(ff, block.feed_forward_output_weight, block.feed_forward_output_bias); + ggml_tensor * normalized = ada_norm(hidden, temb, block); + ggml_tensor * attended = block_index % 2 == 0 + ? attention(normalized, hidden_input, cross_mask_input, block) + : attention(normalized, normalized, nullptr, block); + hidden = ggml_add(graph_context, hidden, attended); + ggml_tensor * ff = ggml_norm(graph_context, hidden, config.ada_norm_epsilon); + ff = linear(ff, block.feed_forward_input_weight, block.feed_forward_input_bias); + ff = ggml_gelu(graph_context, ff); + ff = linear(ff, block.feed_forward_output_weight, block.feed_forward_output_bias); hidden = ggml_add(graph_context, hidden, ff); } - ggml_tensor * output_modulation = linear(ggml_silu(graph_context, temb), - weights.output_modulation_weight, - weights.output_modulation_bias); + ggml_tensor * output_modulation = linear(ggml_silu(graph_context, temb), weights.output_modulation_weight, + weights.output_modulation_bias); // DiT output uses shift then scale, unlike AdaLayerNorm's scale then shift. ggml_tensor * shift = ggml_view_1d(graph_context, output_modulation, width, 0); - ggml_tensor * scale = ggml_view_1d(graph_context, output_modulation, width, - static_cast(width) * sizeof(float)); + ggml_tensor * scale = + ggml_view_1d(graph_context, output_modulation, width, static_cast(width) * sizeof(float)); hidden = ggml_norm(graph_context, hidden, config.output_norm_epsilon); hidden = ggml_mul(graph_context, hidden, ggml_add(graph_context, scale, scalar_one_input)); hidden = ggml_add(graph_context, hidden, shift); hidden = linear(hidden, weights.output_projection_weight, weights.output_projection_bias); - hidden = ggml_relu(graph_context, - linear(hidden, weights.velocity_input_weight, weights.velocity_input_bias)); + hidden = + ggml_relu(graph_context, linear(hidden, weights.velocity_input_weight, weights.velocity_input_bias)); hidden = linear(hidden, weights.velocity_output_weight, weights.velocity_output_bias); - ggml_tensor * velocity = ggml_view_2d( - graph_context, hidden, config.action_dim, config.horizon, hidden->nb[1], - static_cast(config.future_token_count) * hidden->nb[1]); + ggml_tensor * velocity = + ggml_view_2d(graph_context, hidden, config.action_dim, config.horizon, hidden->nb[1], + static_cast(config.future_token_count) * hidden->nb[1]); actions = ggml_add(graph_context, actions, ggml_scale(graph_context, velocity, config.euler_dt)); } @@ -700,12 +545,12 @@ std::unique_ptr GR00TPolicy::load(const std::string & path, int n_t impl->verbosity = verbosity; try { backend_scheduler_config scheduler_config; - scheduler_config.max_nodes = static_cast(kGraphSize); - scheduler_config.parallel = false; + scheduler_config.max_nodes = static_cast(kGraphSize); + scheduler_config.parallel = false; scheduler_config.op_offload = true; backend_loader backend; - if (!backend.load(impl->backend_cpu, impl->backends, impl->scheduler, impl->buft_policy, true, - scheduler_config, verbosity)) { + if (!backend.load(impl->backend_cpu, impl->backends, impl->scheduler, impl->buft_policy, true, scheduler_config, + verbosity)) { error = "failed to initialize StarVLA GR00T backend: " + backend.error(); return nullptr; } @@ -722,11 +567,10 @@ std::unique_ptr GR00TPolicy::load(const std::string & path, int n_t } ggml_backend_buffer_set_usage(impl->loaded.model_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); impl->timestep_table = timestep_projection_table(impl->config); - impl->action_table = action_time_table(impl->config); + impl->action_table = action_time_table(impl->config); if (verbosity >= 1) { - std::fprintf(stderr, - "%s: backend=%s qwen=%d width=%d blocks=%d horizon=%d action_dim=%d profiles=%zu\n", - __func__, mode_name(impl->mode), impl->config.qwen_hidden_dim, impl->config.dit_width, + std::fprintf(stderr, "%s: backend=%s qwen=%d width=%d blocks=%d horizon=%d action_dim=%d profiles=%zu\n", + __func__, backend_mode_name(impl->mode), impl->config.qwen_hidden_dim, impl->config.dit_width, impl->config.block_count, impl->config.horizon, impl->config.action_dim, impl->config.normalization.profiles.size()); } @@ -738,9 +582,8 @@ std::unique_ptr GR00TPolicy::load(const std::string & path, int n_t } bool GR00TPolicy::evaluate(const float * qwen_hidden_states, size_t hidden_element_count, - const uint8_t * qwen_attention_mask, size_t mask_element_count, - const float * initial_noise, size_t noise_element_count, - std::vector & normalized_actions, std::string & error) { + const uint8_t * qwen_attention_mask, size_t mask_element_count, const float * initial_noise, + size_t noise_element_count, std::vector & normalized_actions, std::string & error) { normalized_actions.clear(); error.clear(); if (impl_ == nullptr || impl_->scheduler == nullptr) { @@ -749,8 +592,7 @@ bool GR00TPolicy::evaluate(const float * qwen_hidden_states, size_t hidden_eleme } if (qwen_hidden_states == nullptr || qwen_attention_mask == nullptr || initial_noise == nullptr || mask_element_count == 0 || mask_element_count > static_cast(std::numeric_limits::max()) || - mask_element_count > std::numeric_limits::max() / - static_cast(impl_->config.qwen_hidden_dim) || + mask_element_count > std::numeric_limits::max() / static_cast(impl_->config.qwen_hidden_dim) || hidden_element_count != mask_element_count * static_cast(impl_->config.qwen_hidden_dim)) { error = "StarVLA GR00T Qwen conditioning tensor or attention mask has an incompatible shape"; return false; @@ -789,22 +631,19 @@ bool GR00TPolicy::evaluate(const float * qwen_hidden_states, size_t hidden_eleme return false; } - const int query_count = impl_->config.no_state_sequence_length; + const int query_count = impl_->config.no_state_sequence_length; const int padded_queries = GGML_PAD(query_count, kKQMaskPad); std::vector additive_mask(mask_element_count * static_cast(padded_queries), -std::numeric_limits::infinity()); for (int query = 0; query < query_count; ++query) { float * row = additive_mask.data() + static_cast(query) * mask_element_count; for (size_t token = 0; token < mask_element_count; ++token) { - row[token] = qwen_attention_mask[token] != 0 ? 0.0f : - -std::numeric_limits::infinity(); + row[token] = qwen_attention_mask[token] != 0 ? 0.0f : -std::numeric_limits::infinity(); } } - ggml_backend_tensor_set(impl_->hidden_input, qwen_hidden_states, 0, - hidden_element_count * sizeof(float)); - ggml_backend_tensor_set(impl_->cross_mask_input, additive_mask.data(), 0, - additive_mask.size() * sizeof(float)); + ggml_backend_tensor_set(impl_->hidden_input, qwen_hidden_states, 0, hidden_element_count * sizeof(float)); + ggml_backend_tensor_set(impl_->cross_mask_input, additive_mask.data(), 0, additive_mask.size() * sizeof(float)); ggml_backend_tensor_set(impl_->noise_input, initial_noise, 0, noise_element_count * sizeof(float)); ggml_backend_tensor_set(impl_->timestep_projection_input, impl_->timestep_table.data(), 0, impl_->timestep_table.size() * sizeof(float)); @@ -829,9 +668,8 @@ bool GR00TPolicy::evaluate(const float * qwen_hidden_states, size_t hidden_eleme return true; } -bool GR00TPolicy::unnormalize(const std::vector & normalized_actions, - const std::string & profile_key_value, std::vector & actions, - std::string & error) const { +bool GR00TPolicy::unnormalize(const std::vector & normalized_actions, const std::string & profile_key_value, + std::vector & actions, std::string & error) const { if (impl_ == nullptr) { actions.clear(); error = "StarVLA GR00T policy is not initialized"; @@ -849,7 +687,7 @@ const GR00TPolicyConfig & GR00TPolicy::config() const { } const char * GR00TPolicy::backend_name() const { - return impl_ != nullptr ? mode_name(impl_->mode) : "unknown"; + return impl_ != nullptr ? backend_mode_name(impl_->mode) : "unknown"; } } // namespace robotcpp::starvla diff --git a/src/models/starvla/groot_policy.h b/src/models/starvla/groot_policy.h index c505035..61cc2cc 100644 --- a/src/models/starvla/groot_policy.h +++ b/src/models/starvla/groot_policy.h @@ -15,35 +15,35 @@ struct GR00TPolicyConfig { std::string bundle_uuid; std::string text_filename; std::string mmproj_filename; - int qwen_hidden_dim = 0; + int qwen_hidden_dim = 0; int qwen_input_embedding_dim = 0; - int qwen_vocab_size = 0; + int qwen_vocab_size = 0; std::string cot_template; int image_count = 0; std::vector image_names; int image_processor_min_pixels = 0; int image_processor_max_pixels = 0; - int image_patch_size = 0; - int image_spatial_merge_size = 0; - int image_min_token_count = 0; - int image_max_token_count = 0; - int dit_width = 0; - int block_count = 0; - int attention_head_count = 0; - int attention_head_dim = 0; - int cross_attention_dim = 0; - int feed_forward_dim = 0; - int output_dim = 0; - int mlp_hidden_dim = 0; - int action_dim = 0; - int horizon = 0; - int future_token_count = 0; - int action_position_count = 0; - int no_state_sequence_length = 0; - int timestep_projection_dim = 0; - float ada_norm_epsilon = 0.0f; - float output_norm_epsilon = 0.0f; - float euler_dt = 0.0f; + int image_patch_size = 0; + int image_spatial_merge_size = 0; + int image_min_token_count = 0; + int image_max_token_count = 0; + int dit_width = 0; + int block_count = 0; + int attention_head_count = 0; + int attention_head_dim = 0; + int cross_attention_dim = 0; + int feed_forward_dim = 0; + int output_dim = 0; + int mlp_hidden_dim = 0; + int action_dim = 0; + int horizon = 0; + int future_token_count = 0; + int action_position_count = 0; + int no_state_sequence_length = 0; + int timestep_projection_dim = 0; + float ada_norm_epsilon = 0.0f; + float output_norm_epsilon = 0.0f; + float euler_dt = 0.0f; std::vector timestep_ids; NormalizationConfig normalization; }; @@ -52,7 +52,7 @@ class GR00TPolicy { public: ~GR00TPolicy(); - GR00TPolicy(const GR00TPolicy &) = delete; + GR00TPolicy(const GR00TPolicy &) = delete; GR00TPolicy & operator=(const GR00TPolicy &) = delete; static std::unique_ptr load(const std::string & path, int n_threads, int verbosity, @@ -61,9 +61,8 @@ class GR00TPolicy { // qwen_hidden_states is token-major [token_count, qwen_hidden_dim]. The mask // follows torch SDPA semantics: non-zero entries participate in attention. // initial_noise is token-major [horizon, action_dim]. - bool evaluate(const float * qwen_hidden_states, size_t hidden_element_count, - const uint8_t * qwen_attention_mask, size_t mask_element_count, - const float * initial_noise, size_t noise_element_count, + bool evaluate(const float * qwen_hidden_states, size_t hidden_element_count, const uint8_t * qwen_attention_mask, + size_t mask_element_count, const float * initial_noise, size_t noise_element_count, std::vector & normalized_actions, std::string & error); bool unnormalize(const std::vector & normalized_actions, const std::string & profile_key, std::vector & actions, std::string & error) const; diff --git a/src/models/starvla/groot_prompt.cpp b/src/models/starvla/groot_prompt.cpp index a1bbe78..7b490aa 100644 --- a/src/models/starvla/groot_prompt.cpp +++ b/src/models/starvla/groot_prompt.cpp @@ -6,14 +6,13 @@ namespace robotcpp::starvla { namespace { constexpr const char * kInstructionPlaceholder = "{instruction}"; -constexpr const char * kMtmdMediaMarker = "<__media__>"; +constexpr const char * kMtmdMediaMarker = "<__media__>"; bool contains_nul(const std::string & value) { return value.find('\0') != std::string::npos; } -void replace_all(std::string & value, const std::string & needle, - const std::string & replacement) { +void replace_all(std::string & value, const std::string & needle, const std::string & replacement) { size_t offset = 0; while ((offset = value.find(needle, offset)) != std::string::npos) { value.replace(offset, needle.size(), replacement); @@ -21,9 +20,8 @@ void replace_all(std::string & value, const std::string & needle, } } -bool build_instruction(const char * framework, const std::string & cot_template, - const std::string & task, std::string & instruction, - std::string & error) { +bool build_instruction(const char * framework, const std::string & cot_template, const std::string & task, + std::string & instruction, std::string & error) { instruction.clear(); error.clear(); @@ -32,19 +30,15 @@ bool build_instruction(const char * framework, const std::string & cot_template, return false; } if (contains_nul(task) || contains_nul(cot_template)) { - error = std::string("StarVLA ") + framework + - " prompt contains an embedded NUL byte"; + error = std::string("StarVLA ") + framework + " prompt contains an embedded NUL byte"; return false; } - if (task.find(kMtmdMediaMarker) != std::string::npos || - cot_template.find(kMtmdMediaMarker) != std::string::npos) { - error = std::string("StarVLA ") + framework + - " prompt contains the reserved mtmd media marker"; + if (task.find(kMtmdMediaMarker) != std::string::npos || cot_template.find(kMtmdMediaMarker) != std::string::npos) { + error = std::string("StarVLA ") + framework + " prompt contains the reserved mtmd media marker"; return false; } if (cot_template.find(kInstructionPlaceholder) == std::string::npos) { - error = std::string("StarVLA ") + framework + - " CoT template is missing {instruction}"; + error = std::string("StarVLA ") + framework + " CoT template is missing {instruction}"; return false; } @@ -56,19 +50,18 @@ bool build_instruction(const char * framework, const std::string & cot_template, } // namespace -bool build_groot_instruction(const std::string & cot_template, const std::string & task, - std::string & instruction, std::string & error) { +bool build_groot_instruction(const std::string & cot_template, const std::string & task, std::string & instruction, + std::string & error) { return build_instruction("GR00T", cot_template, task, instruction, error); } -bool build_pi_v3_instruction(const std::string & cot_template, const std::string & task, - std::string & instruction, std::string & error) { +bool build_pi_v3_instruction(const std::string & cot_template, const std::string & task, std::string & instruction, + std::string & error) { return build_instruction("PI_v3", cot_template, task, instruction, error); } -bool build_fast_instruction(const std::string & cot_template, - const std::string & task, - std::string & instruction, std::string & error) { +bool build_fast_instruction(const std::string & cot_template, const std::string & task, std::string & instruction, + std::string & error) { return build_instruction("FAST", cot_template, task, instruction, error); } diff --git a/src/models/starvla/groot_prompt.h b/src/models/starvla/groot_prompt.h index 27e2950..d14eee2 100644 --- a/src/models/starvla/groot_prompt.h +++ b/src/models/starvla/groot_prompt.h @@ -4,14 +4,13 @@ namespace robotcpp::starvla { -bool build_groot_instruction(const std::string & cot_template, const std::string & task, - std::string & instruction, std::string & error); +bool build_groot_instruction(const std::string & cot_template, const std::string & task, std::string & instruction, + std::string & error); -bool build_pi_v3_instruction(const std::string & cot_template, const std::string & task, - std::string & instruction, std::string & error); +bool build_pi_v3_instruction(const std::string & cot_template, const std::string & task, std::string & instruction, + std::string & error); -bool build_fast_instruction(const std::string & cot_template, - const std::string & task, - std::string & instruction, std::string & error); +bool build_fast_instruction(const std::string & cot_template, const std::string & task, std::string & instruction, + std::string & error); } // namespace robotcpp::starvla diff --git a/src/models/starvla/normalization.cpp b/src/models/starvla/normalization.cpp index 9963a6c..8fd03a4 100644 --- a/src/models/starvla/normalization.cpp +++ b/src/models/starvla/normalization.cpp @@ -39,6 +39,10 @@ bool validate_normalization_config(const NormalizationConfig & config, int actio error = "StarVLA policy has no normalization profiles"; return false; } + if (config.default_profile_key.empty()) { + error = "StarVLA policy has no default normalization profile"; + return false; + } std::vector dimension_kind(static_cast(action_dim), 0); for (int32_t dim : config.continuous_dimensions) { @@ -94,6 +98,10 @@ bool validate_normalization_config(const NormalizationConfig & config, int actio } } } + if (std::find(seen_keys.begin(), seen_keys.end(), config.default_profile_key) == seen_keys.end()) { + error = "StarVLA default normalization profile is not present: " + config.default_profile_key; + return false; + } return true; } @@ -101,11 +109,7 @@ const NormalizationProfile * resolve_normalization_profile(const NormalizationCo const std::string & profile_key, std::string & error) { error.clear(); if (profile_key.empty()) { - if (config.profiles.size() == 1) { - return &config.profiles.front(); - } - error = "StarVLA policy has multiple normalization profiles; select one of: " + profile_keys(config); - return nullptr; + return resolve_normalization_profile(config, config.default_profile_key, error); } for (const NormalizationProfile & profile : config.profiles) { if (profile.key == profile_key) { @@ -141,27 +145,22 @@ bool denormalize_actions(const NormalizationConfig & config, const std::string & actions.resize(normalized.size()); for (int step = 0; step < horizon; ++step) { for (int dim = 0; dim < action_dim; ++dim) { - const size_t index = static_cast(step) * static_cast(action_dim) + - static_cast(dim); + const size_t index = static_cast(step) * static_cast(action_dim) + static_cast(dim); const float input_value = normalized[index]; if (!std::isfinite(input_value)) { actions.clear(); error = "StarVLA normalized actions must be finite"; return false; } - const float value = - config.clip_actions ? std::clamp(input_value, -1.0f, 1.0f) - : input_value; + const float value = config.clip_actions ? std::clamp(input_value, -1.0f, 1.0f) : input_value; if (is_binary[static_cast(dim)] != 0) { - const bool active = - config.binary_comparison == "ge" - ? value >= config.binary_threshold - : value > config.binary_threshold; - actions[index] = active ? 1.0f : 0.0f; + const bool active = config.binary_comparison == "ge" ? value >= config.binary_threshold + : value > config.binary_threshold; + actions[index] = active ? 1.0f : 0.0f; } else { - const float low = profile->action_q01[static_cast(dim)]; + const float low = profile->action_q01[static_cast(dim)]; const float high = profile->action_q99[static_cast(dim)]; - actions[index] = (value + 1.0f) * 0.5f * (high - low) + low; + actions[index] = (value + 1.0f) * 0.5f * (high - low) + low; } } } diff --git a/src/models/starvla/normalization.h b/src/models/starvla/normalization.h index 53fe6f1..696830b 100644 --- a/src/models/starvla/normalization.h +++ b/src/models/starvla/normalization.h @@ -14,7 +14,8 @@ struct NormalizationProfile { }; struct NormalizationConfig { - bool clip_actions = false; + std::string default_profile_key; + bool clip_actions = false; float binary_threshold = 0.5f; std::string binary_comparison; std::vector continuous_dimensions; diff --git a/src/models/starvla/oft_image_preprocess.cpp b/src/models/starvla/oft_image_preprocess.cpp index 3192fcd..ecace78 100644 --- a/src/models/starvla/oft_image_preprocess.cpp +++ b/src/models/starvla/oft_image_preprocess.cpp @@ -12,14 +12,14 @@ namespace robotcpp::starvla { namespace { struct RGBImage { - int width = 0; + int width = 0; int height = 0; std::vector pixels; }; struct FilterTable { int kernel_size = 0; - int precision = 0; + int precision = 0; std::vector first; std::vector count; std::vector weights; @@ -27,7 +27,7 @@ struct FilterTable { double keys_cubic(double value) { constexpr double a = -0.5; - value = std::fabs(value); + value = std::fabs(value); if (value < 1.0) { return ((a + 2.0) * value - (a + 3.0)) * value * value + 1.0; } @@ -37,12 +37,10 @@ double keys_cubic(double value) { return 0.0; } -bool validate_resize(const uint8_t * source, int source_width, int source_height, - int source_stride, int target_width, int target_height, - std::string & error) { +bool validate_resize(const uint8_t * source, int source_width, int source_height, int source_stride, int target_width, + int target_height, std::string & error) { error.clear(); - if (source == nullptr || source_width <= 0 || source_height <= 0 || target_width <= 0 || - target_height <= 0) { + if (source == nullptr || source_width <= 0 || source_height <= 0 || target_width <= 0 || target_height <= 0) { error = "StarVLA image resize received an invalid image or dimension"; return false; } @@ -51,8 +49,7 @@ bool validate_resize(const uint8_t * source, int source_width, int source_height error = "StarVLA image stride is smaller than a packed RGB row"; return false; } - const uint64_t output_bytes = static_cast(target_width) * - static_cast(target_height) * 3; + const uint64_t output_bytes = static_cast(target_width) * static_cast(target_height) * 3; if (output_bytes > static_cast(std::numeric_limits::max())) { error = "StarVLA resized image is too large"; return false; @@ -62,9 +59,9 @@ bool validate_resize(const uint8_t * source, int source_width, int source_height RGBImage pack_source(const uint8_t * source, int width, int height, int stride) { RGBImage image; - image.width = width; - image.height = height; - const size_t row_bytes = static_cast(width) * 3; + image.width = width; + image.height = height; + const size_t row_bytes = static_cast(width) * 3; const size_t actual_stride = stride > 0 ? static_cast(stride) : row_bytes; image.pixels.resize(row_bytes * static_cast(height)); for (int row = 0; row < height; ++row) { @@ -76,38 +73,36 @@ RGBImage pack_source(const uint8_t * source, int width, int height, int stride) } FilterTable make_filter_table(int input_size, int output_size, bool pillow_precision) { - const double scale = static_cast(input_size) / static_cast(output_size); + const double scale = static_cast(input_size) / static_cast(output_size); const double filter_scale = std::max(scale, 1.0); - const double support = 2.0 * filter_scale; + const double support = 2.0 * filter_scale; FilterTable table; table.kernel_size = static_cast(std::ceil(support)) * 2 + 1; table.first.resize(static_cast(output_size)); table.count.resize(static_cast(output_size)); - std::vector floating_weights(static_cast(output_size) * - static_cast(table.kernel_size), 0.0); + std::vector floating_weights(static_cast(output_size) * static_cast(table.kernel_size), + 0.0); double maximum_weight = 0.0; for (int output = 0; output < output_size; ++output) { - const double center = (static_cast(output) + 0.5) * scale; - const int first = std::max(static_cast(center - support + 0.5), 0); - const int end = std::min(static_cast(center + support + 0.5), input_size); - const int count = std::max(0, std::min(end - first, table.kernel_size)); + const double center = (static_cast(output) + 0.5) * scale; + const int first = std::max(static_cast(center - support + 0.5), 0); + const int end = std::min(static_cast(center + support + 0.5), input_size); + const int count = std::max(0, std::min(end - first, table.kernel_size)); table.first[static_cast(output)] = first; table.count[static_cast(output)] = count; double sum = 0.0; for (int index = 0; index < count; ++index) { - const double distance = - (static_cast(index + first) - center + 0.5) / filter_scale; - const double weight = keys_cubic(distance); + const double distance = (static_cast(index + first) - center + 0.5) / filter_scale; + const double weight = keys_cubic(distance); floating_weights[static_cast(output) * table.kernel_size + index] = weight; sum += weight; } if (sum != 0.0) { for (int index = 0; index < count; ++index) { - double & weight = - floating_weights[static_cast(output) * table.kernel_size + index]; + double & weight = floating_weights[static_cast(output) * table.kernel_size + index]; weight /= sum; maximum_weight = std::max(maximum_weight, weight); } @@ -118,8 +113,8 @@ FilterTable make_filter_table(int input_size, int output_size, bool pillow_preci table.precision = 22; } else { for (table.precision = 0; table.precision < 22; ++table.precision) { - const int next = static_cast( - 0.5 + maximum_weight * static_cast(uint32_t{1} << (table.precision + 1))); + const int next = + static_cast(0.5 + maximum_weight * static_cast(uint32_t{1} << (table.precision + 1))); if (next >= (1 << 15)) { break; } @@ -129,7 +124,7 @@ FilterTable make_filter_table(int input_size, int output_size, bool pillow_preci const double multiplier = static_cast(uint32_t{1} << table.precision); table.weights.resize(floating_weights.size()); for (size_t index = 0; index < floating_weights.size(); ++index) { - const double scaled = floating_weights[index] * multiplier; + const double scaled = floating_weights[index] * multiplier; table.weights[index] = static_cast(scaled < 0.0 ? scaled - 0.5 : scaled + 0.5); } return table; @@ -140,10 +135,9 @@ uint8_t fixed_point_pixel(int64_t accumulator, int precision) { return static_cast(std::max(0, std::min(255, value))); } -RGBImage resize_horizontal(const RGBImage & source, int target_width, - const FilterTable & table) { +RGBImage resize_horizontal(const RGBImage & source, int target_width, const FilterTable & table) { RGBImage target; - target.width = target_width; + target.width = target_width; target.height = source.height; target.pixels.resize(static_cast(target.width) * target.height * 3); const int64_t rounding = int64_t{1} << (table.precision - 1); @@ -154,14 +148,11 @@ RGBImage resize_horizontal(const RGBImage & source, int target_width, for (int channel = 0; channel < 3; ++channel) { int64_t accumulator = rounding; for (int index = 0; index < count; ++index) { - const size_t source_index = - (static_cast(row) * source.width + first + index) * 3 + channel; - const int32_t weight = - table.weights[static_cast(column) * table.kernel_size + index]; + const size_t source_index = (static_cast(row) * source.width + first + index) * 3 + channel; + const int32_t weight = table.weights[static_cast(column) * table.kernel_size + index]; accumulator += static_cast(source.pixels[source_index]) * weight; } - const size_t target_index = - (static_cast(row) * target.width + column) * 3 + channel; + const size_t target_index = (static_cast(row) * target.width + column) * 3 + channel; target.pixels[target_index] = fixed_point_pixel(accumulator, table.precision); } } @@ -169,10 +160,9 @@ RGBImage resize_horizontal(const RGBImage & source, int target_width, return target; } -RGBImage resize_vertical(const RGBImage & source, int target_height, - const FilterTable & table) { +RGBImage resize_vertical(const RGBImage & source, int target_height, const FilterTable & table) { RGBImage target; - target.width = source.width; + target.width = source.width; target.height = target_height; target.pixels.resize(static_cast(target.width) * target.height * 3); const int64_t rounding = int64_t{1} << (table.precision - 1); @@ -185,12 +175,10 @@ RGBImage resize_vertical(const RGBImage & source, int target_height, for (int index = 0; index < count; ++index) { const size_t source_index = (static_cast(first + index) * source.width + column) * 3 + channel; - const int32_t weight = - table.weights[static_cast(row) * table.kernel_size + index]; + const int32_t weight = table.weights[static_cast(row) * table.kernel_size + index]; accumulator += static_cast(source.pixels[source_index]) * weight; } - const size_t target_index = - (static_cast(row) * target.width + column) * 3 + channel; + const size_t target_index = (static_cast(row) * target.width + column) * 3 + channel; target.pixels[target_index] = fixed_point_pixel(accumulator, table.precision); } } @@ -198,23 +186,21 @@ RGBImage resize_vertical(const RGBImage & source, int target_height, return target; } -bool resize_rgb(const uint8_t * source, int source_width, int source_height, int source_stride, - int target_width, int target_height, bool pillow_precision, - std::vector & target, std::string & error) { +bool resize_rgb(const uint8_t * source, int source_width, int source_height, int source_stride, int target_width, + int target_height, bool pillow_precision, std::vector & target, std::string & error) { target.clear(); - if (!validate_resize(source, source_width, source_height, source_stride, target_width, - target_height, error)) { + if (!validate_resize(source, source_width, source_height, source_stride, target_width, target_height, error)) { return false; } RGBImage current = pack_source(source, source_width, source_height, source_stride); if (source_width != target_width) { - current = resize_horizontal(current, target_width, - make_filter_table(source_width, target_width, pillow_precision)); + current = + resize_horizontal(current, target_width, make_filter_table(source_width, target_width, pillow_precision)); } if (source_height != target_height) { - current = resize_vertical(current, target_height, - make_filter_table(source_height, target_height, pillow_precision)); + current = + resize_vertical(current, target_height, make_filter_table(source_height, target_height, pillow_precision)); } target = std::move(current.pixels); return true; @@ -222,29 +208,26 @@ bool resize_rgb(const uint8_t * source, int source_width, int source_height, int } // namespace -bool resize_pillow_bicubic_rgb(const uint8_t * source, int source_width, int source_height, - int source_stride, int target_width, int target_height, - std::vector & target, std::string & error) { - return resize_rgb(source, source_width, source_height, source_stride, target_width, - target_height, true, target, error); +bool resize_pillow_bicubic_rgb(const uint8_t * source, int source_width, int source_height, int source_stride, + int target_width, int target_height, std::vector & target, + std::string & error) { + return resize_rgb(source, source_width, source_height, source_stride, target_width, target_height, true, target, + error); } -bool resize_torchvision_bicubic_aa_rgb(const uint8_t * source, int source_width, - int source_height, int source_stride, int target_width, - int target_height, std::vector & target, +bool resize_torchvision_bicubic_aa_rgb(const uint8_t * source, int source_width, int source_height, int source_stride, + int target_width, int target_height, std::vector & target, std::string & error) { - return resize_rgb(source, source_width, source_height, source_stride, target_width, - target_height, false, target, error); + return resize_rgb(source, source_width, source_height, source_stride, target_width, target_height, false, target, + error); } -bool qwen3vl_smart_resize_dimensions(int source_width, int source_height, int factor, - int min_pixels, int max_pixels, int & target_width, - int & target_height, std::string & error) { - target_width = 0; +bool qwen3vl_smart_resize_dimensions(int source_width, int source_height, int factor, int min_pixels, int max_pixels, + int & target_width, int & target_height, std::string & error) { + target_width = 0; target_height = 0; error.clear(); - if (source_width <= 0 || source_height <= 0 || factor <= 0 || min_pixels <= 0 || - max_pixels < min_pixels) { + if (source_width <= 0 || source_height <= 0 || factor <= 0 || min_pixels <= 0 || max_pixels < min_pixels) { error = "Qwen3-VL smart resize received an invalid dimension or pixel bound"; return false; } @@ -258,54 +241,49 @@ bool qwen3vl_smart_resize_dimensions(int source_width, int source_height, int fa // Python round() uses ties-to-even. Integer quotient/remainder arithmetic // makes the common first smart_resize step independent of the host FP mode. const auto round_div_ties_to_even = [](int value, int divisor) -> int64_t { - const int64_t quotient = value / divisor; + const int64_t quotient = value / divisor; const int64_t remainder = value % divisor; - const int64_t doubled = remainder * 2; + const int64_t doubled = remainder * 2; if (doubled < divisor || (doubled == divisor && quotient % 2 == 0)) { return quotient; } return quotient + 1; }; - int64_t resized_height = round_div_ties_to_even(source_height, factor) * factor; - int64_t resized_width = round_div_ties_to_even(source_width, factor) * factor; - const int64_t source_pixels = static_cast(source_height) * source_width; + int64_t resized_height = round_div_ties_to_even(source_height, factor) * factor; + int64_t resized_width = round_div_ties_to_even(source_width, factor) * factor; + const int64_t source_pixels = static_cast(source_height) * source_width; const int64_t rounded_pixels = resized_height * resized_width; if (rounded_pixels > max_pixels) { const double beta = std::sqrt(static_cast(source_pixels) / max_pixels); - resized_height = std::max( - factor, static_cast(std::floor(source_height / beta / factor)) * factor); - resized_width = std::max( - factor, static_cast(std::floor(source_width / beta / factor)) * factor); - } else if (rounded_pixels < min_pixels) { - const double beta = std::sqrt(static_cast(min_pixels) / source_pixels); resized_height = - static_cast(std::ceil(source_height * beta / factor)) * factor; + std::max(factor, static_cast(std::floor(source_height / beta / factor)) * factor); resized_width = - static_cast(std::ceil(source_width * beta / factor)) * factor; + std::max(factor, static_cast(std::floor(source_width / beta / factor)) * factor); + } else if (rounded_pixels < min_pixels) { + const double beta = std::sqrt(static_cast(min_pixels) / source_pixels); + resized_height = static_cast(std::ceil(source_height * beta / factor)) * factor; + resized_width = static_cast(std::ceil(source_width * beta / factor)) * factor; } - if (resized_width <= 0 || resized_height <= 0 || - resized_width > std::numeric_limits::max() || + if (resized_width <= 0 || resized_height <= 0 || resized_width > std::numeric_limits::max() || resized_height > std::numeric_limits::max() || resized_width > std::numeric_limits::max() / resized_height) { error = "Qwen3-VL smart resize produced an unsupported output dimension"; return false; } - target_width = static_cast(resized_width); + target_width = static_cast(resized_width); target_height = static_cast(resized_height); return true; } -bool preprocess_qwen3vl_rgb(const uint8_t * source, int source_width, int source_height, - int channels, int source_stride, int patch_size, - int spatial_merge_size, int min_pixels, int max_pixels, - std::vector & target, int & target_width, - int & target_height, int & image_token_count, - std::string & error) { +bool preprocess_qwen3vl_rgb(const uint8_t * source, int source_width, int source_height, int channels, + int source_stride, int patch_size, int spatial_merge_size, int min_pixels, int max_pixels, + std::vector & target, int & target_width, int & target_height, + int & image_token_count, std::string & error) { target.clear(); - target_width = 0; - target_height = 0; + target_width = 0; + target_height = 0; image_token_count = 0; if (channels != 3) { error = "Qwen3-VL input image must be RGB"; @@ -317,33 +295,32 @@ bool preprocess_qwen3vl_rgb(const uint8_t * source, int source_width, int source return false; } const int factor = patch_size * spatial_merge_size; - if (!qwen3vl_smart_resize_dimensions(source_width, source_height, factor, min_pixels, - max_pixels, target_width, target_height, error)) { + if (!qwen3vl_smart_resize_dimensions(source_width, source_height, factor, min_pixels, max_pixels, target_width, + target_height, error)) { return false; } - if (!resize_torchvision_bicubic_aa_rgb(source, source_width, source_height, source_stride, - target_width, target_height, target, error)) { - target_width = 0; + if (!resize_torchvision_bicubic_aa_rgb(source, source_width, source_height, source_stride, target_width, + target_height, target, error)) { + target_width = 0; target_height = 0; return false; } - const int64_t grid_width = target_width / factor; + const int64_t grid_width = target_width / factor; const int64_t grid_height = target_height / factor; - const int64_t tokens = grid_width * grid_height; + const int64_t tokens = grid_width * grid_height; if (tokens <= 0 || tokens > std::numeric_limits::max()) { target.clear(); - target_width = 0; + target_width = 0; target_height = 0; - error = "Qwen3-VL smart resize produced an unsupported image token count"; + error = "Qwen3-VL smart resize produced an unsupported image token count"; return false; } image_token_count = static_cast(tokens); return true; } -bool preprocess_oft_rgb(const uint8_t * source, int source_width, int source_height, - int channels, int source_stride, int training_width, - int training_height, int processor_width, int processor_height, +bool preprocess_oft_rgb(const uint8_t * source, int source_width, int source_height, int channels, int source_stride, + int training_width, int training_height, int processor_width, int processor_height, std::vector & target, std::string & error) { target.clear(); if (channels != 3) { @@ -351,13 +328,12 @@ bool preprocess_oft_rgb(const uint8_t * source, int source_width, int source_hei return false; } std::vector training_image; - if (!resize_pillow_bicubic_rgb(source, source_width, source_height, source_stride, - training_width, training_height, training_image, error)) { + if (!resize_pillow_bicubic_rgb(source, source_width, source_height, source_stride, training_width, training_height, + training_image, error)) { return false; } - return resize_torchvision_bicubic_aa_rgb( - training_image.data(), training_width, training_height, training_width * 3, - processor_width, processor_height, target, error); + return resize_torchvision_bicubic_aa_rgb(training_image.data(), training_width, training_height, training_width * 3, + processor_width, processor_height, target, error); } } // namespace robotcpp::starvla diff --git a/src/models/starvla/oft_image_preprocess.h b/src/models/starvla/oft_image_preprocess.h index e4d1e02..6aaa2a7 100644 --- a/src/models/starvla/oft_image_preprocess.h +++ b/src/models/starvla/oft_image_preprocess.h @@ -6,29 +6,23 @@ namespace robotcpp::starvla { -bool resize_pillow_bicubic_rgb(const uint8_t * source, int source_width, int source_height, - int source_stride, int target_width, int target_height, - std::vector & target, std::string & error); +bool resize_pillow_bicubic_rgb(const uint8_t * source, int source_width, int source_height, int source_stride, + int target_width, int target_height, std::vector & target, std::string & error); -bool resize_torchvision_bicubic_aa_rgb(const uint8_t * source, int source_width, - int source_height, int source_stride, int target_width, - int target_height, std::vector & target, +bool resize_torchvision_bicubic_aa_rgb(const uint8_t * source, int source_width, int source_height, int source_stride, + int target_width, int target_height, std::vector & target, std::string & error); -bool qwen3vl_smart_resize_dimensions(int source_width, int source_height, int factor, - int min_pixels, int max_pixels, int & target_width, - int & target_height, std::string & error); +bool qwen3vl_smart_resize_dimensions(int source_width, int source_height, int factor, int min_pixels, int max_pixels, + int & target_width, int & target_height, std::string & error); -bool preprocess_qwen3vl_rgb(const uint8_t * source, int source_width, int source_height, - int channels, int source_stride, int patch_size, - int spatial_merge_size, int min_pixels, int max_pixels, - std::vector & target, int & target_width, - int & target_height, int & image_token_count, - std::string & error); +bool preprocess_qwen3vl_rgb(const uint8_t * source, int source_width, int source_height, int channels, + int source_stride, int patch_size, int spatial_merge_size, int min_pixels, int max_pixels, + std::vector & target, int & target_width, int & target_height, + int & image_token_count, std::string & error); -bool preprocess_oft_rgb(const uint8_t * source, int source_width, int source_height, - int channels, int source_stride, int training_width, - int training_height, int processor_width, int processor_height, +bool preprocess_oft_rgb(const uint8_t * source, int source_width, int source_height, int channels, int source_stride, + int training_width, int training_height, int processor_width, int processor_height, std::vector & target, std::string & error); } // namespace robotcpp::starvla diff --git a/src/models/starvla/oft_policy.cpp b/src/models/starvla/oft_policy.cpp index 48bbd62..3f8bccc 100644 --- a/src/models/starvla/oft_policy.cpp +++ b/src/models/starvla/oft_policy.cpp @@ -5,10 +5,10 @@ #include "gguf.h" #include "models/ggml_backend.h" #include "models/gguf_loader.h" +#include "models/starvla/policy_gguf.h" #include #include -#include #include #include #include @@ -19,138 +19,30 @@ namespace robotcpp::starvla { namespace { struct OFTBlockWeights { - ggml_tensor * norm_weight = nullptr; - ggml_tensor * norm_bias = nullptr; + ggml_tensor * norm_weight = nullptr; + ggml_tensor * norm_bias = nullptr; ggml_tensor * linear_weight = nullptr; - ggml_tensor * linear_bias = nullptr; + ggml_tensor * linear_bias = nullptr; }; struct OFTWeights { ggml_tensor * input_norm_weight = nullptr; - ggml_tensor * input_norm_bias = nullptr; + ggml_tensor * input_norm_bias = nullptr; ggml_tensor * input_proj_weight = nullptr; - ggml_tensor * input_proj_bias = nullptr; + ggml_tensor * input_proj_bias = nullptr; std::vector blocks; ggml_tensor * output_norm_weight = nullptr; - ggml_tensor * output_norm_bias = nullptr; + ggml_tensor * output_norm_bias = nullptr; ggml_tensor * output_proj_weight = nullptr; - ggml_tensor * output_proj_bias = nullptr; + ggml_tensor * output_proj_bias = nullptr; }; -int require_key(gguf_context * gguf, const char * key, gguf_type type) { - const int index = gguf_find_key(gguf, key); - if (index < 0) { - throw std::runtime_error(std::string("missing required StarVLA GGUF metadata: ") + key); - } - if (gguf_get_kv_type(gguf, index) != type) { - throw std::runtime_error(std::string("invalid StarVLA GGUF metadata type: ") + key); - } - return index; -} - -std::string require_string(gguf_context * gguf, const char * key) { - return gguf_get_val_str(gguf, require_key(gguf, key, GGUF_TYPE_STRING)); -} - -int require_i32(gguf_context * gguf, const char * key) { - return gguf_get_val_i32(gguf, require_key(gguf, key, GGUF_TYPE_INT32)); -} - -float require_f32(gguf_context * gguf, const char * key) { - return gguf_get_val_f32(gguf, require_key(gguf, key, GGUF_TYPE_FLOAT32)); -} - -bool require_bool(gguf_context * gguf, const char * key) { - return gguf_get_val_bool(gguf, require_key(gguf, key, GGUF_TYPE_BOOL)); -} - -int require_array(gguf_context * gguf, const char * key, gguf_type element_type) { - const int index = require_key(gguf, key, GGUF_TYPE_ARRAY); - if (gguf_get_arr_type(gguf, index) != element_type) { - throw std::runtime_error(std::string("invalid StarVLA GGUF array element type: ") + key); - } - return index; -} - -std::vector require_string_array(gguf_context * gguf, const char * key) { - const int index = require_array(gguf, key, GGUF_TYPE_STRING); - const size_t count = gguf_get_arr_n(gguf, index); - std::vector values; - values.reserve(count); - for (size_t i = 0; i < count; ++i) { - values.emplace_back(gguf_get_arr_str(gguf, index, i)); - } - return values; -} - -std::vector require_i32_array(gguf_context * gguf, const char * key) { - const int index = require_array(gguf, key, GGUF_TYPE_INT32); - const size_t count = gguf_get_arr_n(gguf, index); - if (count == 0) { - return {}; - } - const auto * data = static_cast(gguf_get_arr_data(gguf, index)); - if (data == nullptr) { - throw std::runtime_error(std::string("missing StarVLA GGUF array data: ") + key); - } - return std::vector(data, data + count); -} - -std::vector require_f32_array(gguf_context * gguf, const char * key) { - const int index = require_array(gguf, key, GGUF_TYPE_FLOAT32); - const size_t count = gguf_get_arr_n(gguf, index); - if (count == 0) { - return {}; - } - const auto * data = static_cast(gguf_get_arr_data(gguf, index)); - if (data == nullptr) { - throw std::runtime_error(std::string("missing StarVLA GGUF array data: ") + key); - } - return std::vector(data, data + count); -} - -std::vector require_bool_array(gguf_context * gguf, const char * key) { - const int index = require_array(gguf, key, GGUF_TYPE_BOOL); - const size_t count = gguf_get_arr_n(gguf, index); - const auto * data = static_cast(gguf_get_arr_data(gguf, index)); - if (data == nullptr && count != 0) { - throw std::runtime_error(std::string("missing StarVLA GGUF array data: ") + key); - } - std::vector values(count); - for (size_t i = 0; i < count; ++i) { - values[i] = data[i] != 0 ? 1 : 0; - } - return values; -} - -std::string profile_key(int profile_index, const char * suffix) { - return "starvla.normalization.profile." + std::to_string(profile_index) + "." + suffix; -} - -bool has_shape(const ggml_tensor * tensor, std::initializer_list expected) { - if (tensor == nullptr || static_cast(ggml_n_dims(tensor)) != expected.size()) { - return false; - } - size_t dimension = 0; - for (int64_t value : expected) { - if (tensor->ne[dimension++] != value) { - return false; - } - } - return true; -} - -const char * mode_name(backend_mode mode) { - switch (mode) { - case backend_mode::cpu: - return "cpu"; - case backend_mode::cuda: - return "cuda"; - case backend_mode::metal: - return "metal"; - } - return "unknown"; -} +using detail::has_shape; +using detail::require_bool; +using detail::require_f32; +using detail::require_i32; +using detail::require_string; +using detail::require_string_array; class OFTGGUFLoader final : public gguf_loader { public: @@ -162,140 +54,92 @@ class OFTGGUFLoader final : public gguf_loader { if (architecture != "starvla-policy") { throw std::runtime_error("StarVLA policy GGUF has incompatible general.architecture: " + architecture); } - if (require_i32(gguf, "starvla.schema_version") != 1 || - require_string(gguf, "starvla.framework") != "oft") { + if (require_i32(gguf, "starvla.schema_version") != 1 || require_string(gguf, "starvla.framework") != "oft") { throw std::runtime_error("StarVLA policy GGUF is not a supported Qwen OFT schema"); } config_.backbone_arch = require_string(gguf, "starvla.backbone.arch"); - if (config_.backbone_arch != "qwen3_vl" && - config_.backbone_arch != "qwen2_5_vl") { - throw std::runtime_error( - "StarVLA OFT policy has an unsupported Qwen backbone"); + if (config_.backbone_arch != "qwen3_vl" && config_.backbone_arch != "qwen2_5_vl") { + throw std::runtime_error("StarVLA OFT policy has an unsupported Qwen backbone"); } config_.bundle_uuid = require_string(gguf, "starvla.bundle.uuid"); if (config_.bundle_uuid.empty()) { throw std::runtime_error("StarVLA policy bundle UUID is missing"); } - config_.text_filename = require_string(gguf, "starvla.component.text.filename"); + config_.text_filename = require_string(gguf, "starvla.component.text.filename"); config_.mmproj_filename = require_string(gguf, "starvla.component.mmproj.filename"); if (config_.text_filename.empty() || config_.mmproj_filename.empty()) { throw std::runtime_error("StarVLA policy component filenames must be non-empty"); } - config_.input_dim = require_i32(gguf, "starvla.qwen.hidden_size"); + config_.input_dim = require_i32(gguf, "starvla.qwen.hidden_size"); config_.input_embedding_dim = require_i32(gguf, "starvla.qwen.input_embedding_size"); - config_.vocab_size = require_i32(gguf, "starvla.qwen.vocab_size"); - config_.hidden_dim = require_i32(gguf, "starvla.oft.hidden_size"); - config_.block_count = require_i32(gguf, "starvla.oft.block_count"); - config_.action_dim = require_i32(gguf, "starvla.action.dimension"); - config_.horizon = require_i32(gguf, "starvla.action.horizon"); - config_.layer_norm_epsilon = require_f32(gguf, "starvla.oft.layer_norm_epsilon"); - if (config_.input_dim <= 0 || config_.input_embedding_dim <= 0 || - config_.vocab_size <= 0 || config_.hidden_dim <= 0 || - config_.block_count <= 0 || - config_.action_dim <= 0 || config_.horizon <= 0 || !std::isfinite(config_.layer_norm_epsilon) || - config_.layer_norm_epsilon <= 0.0f) { + config_.vocab_size = require_i32(gguf, "starvla.qwen.vocab_size"); + config_.hidden_dim = require_i32(gguf, "starvla.oft.hidden_size"); + config_.block_count = require_i32(gguf, "starvla.oft.block_count"); + config_.action_dim = require_i32(gguf, "starvla.action.dimension"); + config_.horizon = require_i32(gguf, "starvla.action.horizon"); + config_.layer_norm_epsilon = require_f32(gguf, "starvla.oft.layer_norm_epsilon"); + if (config_.input_dim <= 0 || config_.input_embedding_dim <= 0 || config_.vocab_size <= 0 || + config_.hidden_dim <= 0 || config_.block_count <= 0 || config_.action_dim <= 0 || config_.horizon <= 0 || + !std::isfinite(config_.layer_norm_epsilon) || config_.layer_norm_epsilon <= 0.0f) { throw std::runtime_error("StarVLA OFT policy metadata has incompatible dimensions"); } - config_.prompt.horizon = config_.horizon; - config_.prompt.action_token = require_string(gguf, "starvla.prompt.action_token"); + config_.prompt.horizon = config_.horizon; + config_.prompt.action_token = require_string(gguf, "starvla.prompt.action_token"); config_.prompt.action_suffix = require_string(gguf, "starvla.prompt.action_suffix"); - config_.prompt.cot_enabled = require_bool(gguf, "starvla.prompt.cot_enabled"); - config_.prompt.cot_template = require_string(gguf, "starvla.prompt.cot_template"); - config_.prompt.state_bins = require_i32(gguf, "starvla.prompt.state_bins"); + config_.prompt.cot_enabled = require_bool(gguf, "starvla.prompt.cot_enabled"); + config_.prompt.cot_template = require_string(gguf, "starvla.prompt.cot_template"); + config_.prompt.state_bins = require_i32(gguf, "starvla.prompt.state_bins"); config_.prompt.state_bin_min = require_f32(gguf, "starvla.prompt.state_bin_min"); config_.prompt.state_bin_max = require_f32(gguf, "starvla.prompt.state_bin_max"); - config_.prompt.state_clip = require_bool(gguf, "starvla.prompt.state_clip"); - config_.action_token_id = require_i32(gguf, "starvla.prompt.action_token_id"); + config_.prompt.state_clip = require_bool(gguf, "starvla.prompt.state_clip"); + config_.action_token_id = require_i32(gguf, "starvla.prompt.action_token_id"); std::string prompt_error; if (!validate_oft_prompt_config(config_.prompt, prompt_error) || config_.action_token_id < 0) { - throw std::runtime_error(prompt_error.empty() ? - "StarVLA OFT token/template metadata is incompatible" : - prompt_error); + throw std::runtime_error(prompt_error.empty() ? "StarVLA OFT token/template metadata is incompatible" + : prompt_error); } - config_.image_count = require_i32(gguf, "starvla.image.count"); - config_.image_names = require_string_array(gguf, "starvla.image.names"); - config_.image_processor_min_pixels = - require_i32(gguf, "starvla.image.processor_min_pixels"); - config_.image_processor_max_pixels = - require_i32(gguf, "starvla.image.processor_max_pixels"); - config_.image_patch_size = require_i32(gguf, "starvla.image.patch_size"); - config_.image_spatial_merge_size = - require_i32(gguf, "starvla.image.spatial_merge_size"); - config_.image_min_token_count = - require_i32(gguf, "starvla.image.min_token_count"); - config_.image_max_token_count = - require_i32(gguf, "starvla.image.max_token_count"); - if (config_.image_count <= 0 || - config_.image_names.size() != static_cast(config_.image_count) || + config_.image_count = require_i32(gguf, "starvla.image.count"); + config_.image_names = require_string_array(gguf, "starvla.image.names"); + config_.image_processor_min_pixels = require_i32(gguf, "starvla.image.processor_min_pixels"); + config_.image_processor_max_pixels = require_i32(gguf, "starvla.image.processor_max_pixels"); + config_.image_patch_size = require_i32(gguf, "starvla.image.patch_size"); + config_.image_spatial_merge_size = require_i32(gguf, "starvla.image.spatial_merge_size"); + config_.image_min_token_count = require_i32(gguf, "starvla.image.min_token_count"); + config_.image_max_token_count = require_i32(gguf, "starvla.image.max_token_count"); + if (config_.image_count <= 0 || config_.image_names.size() != static_cast(config_.image_count) || config_.image_processor_min_pixels <= 0 || - config_.image_processor_max_pixels < config_.image_processor_min_pixels || - config_.image_patch_size <= 0 || - config_.image_spatial_merge_size <= 0 || - config_.image_min_token_count <= 0 || + config_.image_processor_max_pixels < config_.image_processor_min_pixels || config_.image_patch_size <= 0 || + config_.image_spatial_merge_size <= 0 || config_.image_min_token_count <= 0 || config_.image_max_token_count < config_.image_min_token_count) { throw std::runtime_error("StarVLA OFT image metadata is incompatible"); } - NormalizationConfig & normalization = config_.normalization; - normalization.clip_actions = require_bool(gguf, "starvla.normalization.clip_actions"); - normalization.binary_threshold = require_f32(gguf, "starvla.normalization.binary_threshold"); - normalization.binary_comparison = require_string(gguf, "starvla.normalization.binary_comparison"); - normalization.continuous_dimensions = - require_i32_array(gguf, "starvla.action.continuous_dimensions"); - normalization.binary_dimensions = require_i32_array(gguf, "starvla.action.binary_dimensions"); - - const int profile_count = require_i32(gguf, "starvla.normalization.profile_count"); - const std::vector keys = - require_string_array(gguf, "starvla.normalization.profile_keys"); - if (profile_count <= 0 || keys.size() != static_cast(profile_count)) { - throw std::runtime_error("StarVLA normalization profile count is inconsistent"); - } - normalization.profiles.clear(); - normalization.profiles.reserve(static_cast(profile_count)); - for (int i = 0; i < profile_count; ++i) { - NormalizationProfile profile; - const std::string key_field = profile_key(i, "key"); - const std::string q01_field = profile_key(i, "action_q01"); - const std::string q99_field = profile_key(i, "action_q99"); - const std::string mask_field = profile_key(i, "action_mask"); - profile.key = require_string(gguf, key_field.c_str()); - profile.action_q01 = require_f32_array(gguf, q01_field.c_str()); - profile.action_q99 = require_f32_array(gguf, q99_field.c_str()); - profile.action_mask = require_bool_array(gguf, mask_field.c_str()); - if (profile.key != keys[static_cast(i)]) { - throw std::runtime_error("StarVLA normalization profile order is inconsistent"); - } - normalization.profiles.push_back(std::move(profile)); - } - std::string normalization_error; - if (!validate_normalization_config(normalization, config_.action_dim, normalization_error)) { - throw std::runtime_error(normalization_error); - } + config_.normalization = detail::require_normalization(gguf, config_.action_dim); return true; } bool bind_tensors(ggml_context * ctx_data) override { weights_.input_norm_weight = require_tensor(ctx_data, "starvla.policy.oft.input_norm.weight"); - weights_.input_norm_bias = require_tensor(ctx_data, "starvla.policy.oft.input_norm.bias"); + weights_.input_norm_bias = require_tensor(ctx_data, "starvla.policy.oft.input_norm.bias"); weights_.input_proj_weight = require_tensor(ctx_data, "starvla.policy.oft.input_proj.weight"); - weights_.input_proj_bias = require_tensor(ctx_data, "starvla.policy.oft.input_proj.bias"); + weights_.input_proj_bias = require_tensor(ctx_data, "starvla.policy.oft.input_proj.bias"); weights_.blocks.clear(); weights_.blocks.reserve(static_cast(config_.block_count)); for (int block = 0; block < config_.block_count; ++block) { const std::string prefix = "starvla.policy.oft.block." + std::to_string(block) + "."; OFTBlockWeights current; - current.norm_weight = require_tensor(ctx_data, prefix + "norm.weight"); - current.norm_bias = require_tensor(ctx_data, prefix + "norm.bias"); + current.norm_weight = require_tensor(ctx_data, prefix + "norm.weight"); + current.norm_bias = require_tensor(ctx_data, prefix + "norm.bias"); current.linear_weight = require_tensor(ctx_data, prefix + "linear.weight"); - current.linear_bias = require_tensor(ctx_data, prefix + "linear.bias"); + current.linear_bias = require_tensor(ctx_data, prefix + "linear.bias"); weights_.blocks.push_back(current); } weights_.output_norm_weight = require_tensor(ctx_data, "starvla.policy.oft.output_norm.weight"); - weights_.output_norm_bias = require_tensor(ctx_data, "starvla.policy.oft.output_norm.bias"); + weights_.output_norm_bias = require_tensor(ctx_data, "starvla.policy.oft.output_norm.bias"); weights_.output_proj_weight = require_tensor(ctx_data, "starvla.policy.oft.output_proj.weight"); - weights_.output_proj_bias = require_tensor(ctx_data, "starvla.policy.oft.output_proj.bias"); + weights_.output_proj_bias = require_tensor(ctx_data, "starvla.policy.oft.output_proj.bias"); if (!has_shape(weights_.input_norm_weight, {config_.input_dim}) || !has_shape(weights_.input_norm_bias, {config_.input_dim}) || @@ -333,13 +177,13 @@ struct OFTPolicy::Impl { std::vector backends; ggml_backend_sched_t scheduler = nullptr; backend_buft_policy buft_policy; - backend_mode mode = backend_mode::cpu; - int n_threads = 0; - int verbosity = 0; + backend_mode mode = backend_mode::cpu; + int n_threads = 0; + int verbosity = 0; ggml_context * graph_context = nullptr; - ggml_cgraph * graph = nullptr; - ggml_tensor * input = nullptr; - ggml_tensor * output = nullptr; + ggml_cgraph * graph = nullptr; + ggml_tensor * input = nullptr; + ggml_tensor * output = nullptr; ~Impl() { if (scheduler != nullptr) { @@ -375,10 +219,10 @@ struct OFTPolicy::Impl { void build_graph() { const size_t graph_size = GGML_DEFAULT_GRAPH_SIZE; ggml_init_params params{}; - params.mem_size = graph_size * ggml_tensor_overhead() + ggml_graph_overhead_custom(graph_size, false); + params.mem_size = graph_size * ggml_tensor_overhead() + ggml_graph_overhead_custom(graph_size, false); params.mem_buffer = nullptr; - params.no_alloc = true; - graph_context = ggml_init(params); + params.no_alloc = true; + graph_context = ggml_init(params); if (graph_context == nullptr) { throw std::runtime_error("failed to initialize StarVLA OFT graph context"); } @@ -388,7 +232,7 @@ struct OFTPolicy::Impl { }; auto layer_norm = [&](ggml_tensor * value, ggml_tensor * weight, ggml_tensor * bias) { ggml_tensor * normalized = ggml_norm(graph_context, value, config.layer_norm_epsilon); - normalized = ggml_mul(graph_context, normalized, f32_vector(weight)); + normalized = ggml_mul(graph_context, normalized, f32_vector(weight)); return ggml_add(graph_context, normalized, f32_vector(bias)); }; auto linear = [&](ggml_tensor * value, ggml_tensor * weight, ggml_tensor * bias) { @@ -405,12 +249,12 @@ struct OFTPolicy::Impl { current = ggml_relu(graph_context, linear(current, weights.input_proj_weight, weights.input_proj_bias)); for (const OFTBlockWeights & block : weights.blocks) { ggml_tensor * residual = current; - current = layer_norm(current, block.norm_weight, block.norm_bias); - current = ggml_relu(graph_context, linear(current, block.linear_weight, block.linear_bias)); - current = ggml_add(graph_context, current, residual); + current = layer_norm(current, block.norm_weight, block.norm_bias); + current = ggml_relu(graph_context, linear(current, block.linear_weight, block.linear_bias)); + current = ggml_add(graph_context, current, residual); } current = layer_norm(current, weights.output_norm_weight, weights.output_norm_bias); - output = linear(current, weights.output_proj_weight, weights.output_proj_bias); + output = linear(current, weights.output_proj_weight, weights.output_proj_bias); ggml_set_name(output, "starvla_oft_normalized_actions"); ggml_set_output(output); @@ -443,12 +287,12 @@ std::unique_ptr OFTPolicy::load(const std::string & path, int n_threa impl->verbosity = verbosity; try { backend_scheduler_config scheduler_config; - scheduler_config.max_nodes = GGML_DEFAULT_GRAPH_SIZE; - scheduler_config.parallel = false; + scheduler_config.max_nodes = GGML_DEFAULT_GRAPH_SIZE; + scheduler_config.parallel = false; scheduler_config.op_offload = true; backend_loader backend; - if (!backend.load(impl->backend_cpu, impl->backends, impl->scheduler, impl->buft_policy, true, - scheduler_config, verbosity)) { + if (!backend.load(impl->backend_cpu, impl->backends, impl->scheduler, impl->buft_policy, true, scheduler_config, + verbosity)) { error = "failed to initialize StarVLA OFT backend: " + backend.error(); return nullptr; } @@ -466,9 +310,8 @@ std::unique_ptr OFTPolicy::load(const std::string & path, int n_threa ggml_backend_buffer_set_usage(impl->loaded.model_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); impl->build_graph(); if (verbosity >= 1) { - std::fprintf(stderr, - "%s: backend=%s input=%d hidden=%d blocks=%d horizon=%d action_dim=%d profiles=%zu\n", - __func__, mode_name(impl->mode), impl->config.input_dim, impl->config.hidden_dim, + std::fprintf(stderr, "%s: backend=%s input=%d hidden=%d blocks=%d horizon=%d action_dim=%d profiles=%zu\n", + __func__, backend_mode_name(impl->mode), impl->config.input_dim, impl->config.hidden_dim, impl->config.block_count, impl->config.horizon, impl->config.action_dim, impl->config.normalization.profiles.size()); } @@ -479,23 +322,21 @@ std::unique_ptr OFTPolicy::load(const std::string & path, int n_threa return std::unique_ptr(new OFTPolicy(std::move(impl))); } -bool OFTPolicy::evaluate(const float * action_queries, size_t element_count, - std::vector & normalized_actions, std::string & error) { +bool OFTPolicy::evaluate(const float * action_queries, size_t element_count, std::vector & normalized_actions, + std::string & error) { normalized_actions.clear(); error.clear(); if (impl_ == nullptr) { error = "StarVLA OFT policy is not initialized"; return false; } - const size_t expected = static_cast(impl_->config.horizon) * - static_cast(impl_->config.input_dim); + const size_t expected = static_cast(impl_->config.horizon) * static_cast(impl_->config.input_dim); if (action_queries == nullptr || element_count != expected) { error = "StarVLA OFT action-query tensor has an incompatible shape"; return false; } - if (impl_->scheduler == nullptr || impl_->graph == nullptr || impl_->input == nullptr || - impl_->output == nullptr) { + if (impl_->scheduler == nullptr || impl_->graph == nullptr || impl_->input == nullptr || impl_->output == nullptr) { error = "StarVLA OFT policy graph is not initialized"; return false; } @@ -507,8 +348,8 @@ bool OFTPolicy::evaluate(const float * action_queries, size_t element_count, return false; } - const size_t output_count = static_cast(impl_->config.horizon) * - static_cast(impl_->config.action_dim); + const size_t output_count = + static_cast(impl_->config.horizon) * static_cast(impl_->config.action_dim); normalized_actions.resize(output_count); ggml_backend_tensor_get(impl_->output, normalized_actions.data(), 0, output_count * sizeof(float)); return true; @@ -533,7 +374,7 @@ const OFTPolicyConfig & OFTPolicy::config() const { } const char * OFTPolicy::backend_name() const { - return impl_ != nullptr ? mode_name(impl_->mode) : "unknown"; + return impl_ != nullptr ? backend_mode_name(impl_->mode) : "unknown"; } } // namespace robotcpp::starvla diff --git a/src/models/starvla/oft_policy.h b/src/models/starvla/oft_policy.h index 44a1aed..8416e1c 100644 --- a/src/models/starvla/oft_policy.h +++ b/src/models/starvla/oft_policy.h @@ -15,24 +15,24 @@ struct OFTPolicyConfig { std::string bundle_uuid; std::string text_filename; std::string mmproj_filename; - int input_dim = 0; - int input_embedding_dim = 0; - int vocab_size = 0; - int hidden_dim = 0; - int block_count = 0; - int action_dim = 0; - int horizon = 0; + int input_dim = 0; + int input_embedding_dim = 0; + int vocab_size = 0; + int hidden_dim = 0; + int block_count = 0; + int action_dim = 0; + int horizon = 0; float layer_norm_epsilon = 0.0f; OFTPromptConfig prompt; int action_token_id = 0; - int image_count = 0; + int image_count = 0; std::vector image_names; int image_processor_min_pixels = 0; int image_processor_max_pixels = 0; - int image_patch_size = 0; - int image_spatial_merge_size = 0; - int image_min_token_count = 0; - int image_max_token_count = 0; + int image_patch_size = 0; + int image_spatial_merge_size = 0; + int image_min_token_count = 0; + int image_max_token_count = 0; NormalizationConfig normalization; }; @@ -40,11 +40,10 @@ class OFTPolicy { public: ~OFTPolicy(); - OFTPolicy(const OFTPolicy &) = delete; + OFTPolicy(const OFTPolicy &) = delete; OFTPolicy & operator=(const OFTPolicy &) = delete; - static std::unique_ptr load(const std::string & path, int n_threads, int verbosity, - std::string & error); + static std::unique_ptr load(const std::string & path, int n_threads, int verbosity, std::string & error); bool evaluate(const float * action_queries, size_t element_count, std::vector & normalized_actions, std::string & error); diff --git a/src/models/starvla/oft_prompt.cpp b/src/models/starvla/oft_prompt.cpp index 28fb174..a2fc1da 100644 --- a/src/models/starvla/oft_prompt.cpp +++ b/src/models/starvla/oft_prompt.cpp @@ -9,7 +9,7 @@ namespace robotcpp::starvla { namespace { constexpr const char * kInstructionPlaceholder = "{instruction}"; -constexpr const char * kMtmdMediaMarker = "<__media__>"; +constexpr const char * kMtmdMediaMarker = "<__media__>"; std::string repeat(const std::string & value, int count) { std::string result; @@ -28,8 +28,8 @@ void replace_all(std::string & value, const std::string & needle, const std::str } } -bool discretize_state(const OFTPromptConfig & config, const std::vector & state, - std::string & output, std::string & error) { +bool discretize_state(const OFTPromptConfig & config, const std::vector & state, std::string & output, + std::string & error) { if (state.empty()) { output.clear(); return true; @@ -38,7 +38,7 @@ bool discretize_state(const OFTPromptConfig & config, const std::vector & std::ostringstream stream; const double minimum = static_cast(config.state_bin_min); const double maximum = static_cast(config.state_bin_max); - const double step = (maximum - minimum) / static_cast(config.state_bins); + const double step = (maximum - minimum) / static_cast(config.state_bins); for (size_t i = 0; i < state.size(); ++i) { double value = static_cast(state[i]); if (!std::isfinite(value)) { @@ -77,8 +77,8 @@ bool validate_oft_prompt_config(const OFTPromptConfig & config, std::string & er return false; } const std::string expected_suffix = " Please predict the next " + std::to_string(config.horizon) + - " robot actions: " + - repeat(config.action_token, config.horizon) + "."; + " robot actions: " + repeat(config.action_token, config.horizon) + + "."; if (config.action_suffix != expected_suffix) { error = "StarVLA OFT action suffix does not match its horizon/token contract"; return false; @@ -87,8 +87,7 @@ bool validate_oft_prompt_config(const OFTPromptConfig & config, std::string & er error = "StarVLA OFT CoT template is missing {instruction}"; return false; } - if (config.state_bins <= 0 || - !std::isfinite(config.state_bin_min) || !std::isfinite(config.state_bin_max) || + if (config.state_bins <= 0 || !std::isfinite(config.state_bin_min) || !std::isfinite(config.state_bin_max) || config.state_bin_max <= config.state_bin_min) { error = "StarVLA OFT state prompt metadata is incompatible"; return false; @@ -96,9 +95,8 @@ bool validate_oft_prompt_config(const OFTPromptConfig & config, std::string & er return true; } -bool build_oft_instruction(const OFTPromptConfig & config, const std::string & task, - const std::vector & state, std::string & instruction, - std::string & error) { +bool build_oft_instruction(const OFTPromptConfig & config, const std::string & task, const std::vector & state, + std::string & instruction, std::string & error) { instruction.clear(); if (!validate_oft_prompt_config(config, error)) { return false; @@ -127,8 +125,7 @@ bool build_oft_instruction(const OFTPromptConfig & config, const std::string & t return true; } -std::string build_qwen_media_content(size_t image_count, const std::string & instruction, - const char * media_marker) { +std::string build_qwen_media_content(size_t image_count, const std::string & instruction, const char * media_marker) { const std::string marker = media_marker == nullptr ? std::string() : std::string(media_marker); std::string content; content.reserve(marker.size() * image_count + instruction.size()); @@ -139,9 +136,8 @@ std::string build_qwen_media_content(size_t image_count, const std::string & ins return content; } -bool find_last_token_positions(const std::vector & token_ids, int32_t token_id, - size_t count, std::vector & positions, - std::string & error) { +bool find_last_token_positions(const std::vector & token_ids, int32_t token_id, size_t count, + std::vector & positions, std::string & error) { positions.clear(); error.clear(); if (count == 0) { diff --git a/src/models/starvla/oft_prompt.h b/src/models/starvla/oft_prompt.h index 4f70435..393a7ea 100644 --- a/src/models/starvla/oft_prompt.h +++ b/src/models/starvla/oft_prompt.h @@ -13,23 +13,20 @@ struct OFTPromptConfig { std::string action_suffix; bool cot_enabled = false; std::string cot_template; - int state_bins = 0; + int state_bins = 0; float state_bin_min = 0.0f; float state_bin_max = 0.0f; - bool state_clip = false; + bool state_clip = false; }; bool validate_oft_prompt_config(const OFTPromptConfig & config, std::string & error); -bool build_oft_instruction(const OFTPromptConfig & config, const std::string & task, - const std::vector & state, std::string & instruction, - std::string & error); +bool build_oft_instruction(const OFTPromptConfig & config, const std::string & task, const std::vector & state, + std::string & instruction, std::string & error); -std::string build_qwen_media_content(size_t image_count, const std::string & instruction, - const char * media_marker); +std::string build_qwen_media_content(size_t image_count, const std::string & instruction, const char * media_marker); -bool find_last_token_positions(const std::vector & token_ids, int32_t token_id, - size_t count, std::vector & positions, - std::string & error); +bool find_last_token_positions(const std::vector & token_ids, int32_t token_id, size_t count, + std::vector & positions, std::string & error); } // namespace robotcpp::starvla diff --git a/src/models/starvla/pi_policy.cpp b/src/models/starvla/pi_policy.cpp index 26fcc46..70e24c7 100644 --- a/src/models/starvla/pi_policy.cpp +++ b/src/models/starvla/pi_policy.cpp @@ -5,11 +5,11 @@ #include "gguf.h" #include "models/ggml_backend.h" #include "models/gguf_loader.h" +#include "models/starvla/policy_gguf.h" #include #include #include -#include #include #include #include @@ -22,171 +22,54 @@ namespace { constexpr size_t kGraphSize = 16384; struct PIBlockWeights { - ggml_tensor * ada_norm_weight = nullptr; - ggml_tensor * ada_norm_bias = nullptr; - ggml_tensor * query_weight = nullptr; - ggml_tensor * query_bias = nullptr; - ggml_tensor * key_weight = nullptr; - ggml_tensor * key_bias = nullptr; - ggml_tensor * value_weight = nullptr; - ggml_tensor * value_bias = nullptr; - ggml_tensor * attention_output_weight = nullptr; - ggml_tensor * attention_output_bias = nullptr; - ggml_tensor * feed_forward_input_weight = nullptr; - ggml_tensor * feed_forward_input_bias = nullptr; + ggml_tensor * ada_norm_weight = nullptr; + ggml_tensor * ada_norm_bias = nullptr; + ggml_tensor * query_weight = nullptr; + ggml_tensor * query_bias = nullptr; + ggml_tensor * key_weight = nullptr; + ggml_tensor * key_bias = nullptr; + ggml_tensor * value_weight = nullptr; + ggml_tensor * value_bias = nullptr; + ggml_tensor * attention_output_weight = nullptr; + ggml_tensor * attention_output_bias = nullptr; + ggml_tensor * feed_forward_input_weight = nullptr; + ggml_tensor * feed_forward_input_bias = nullptr; ggml_tensor * feed_forward_output_weight = nullptr; - ggml_tensor * feed_forward_output_bias = nullptr; + ggml_tensor * feed_forward_output_bias = nullptr; }; struct PIWeights { - ggml_tensor * timestep_input_weight = nullptr; - ggml_tensor * timestep_input_bias = nullptr; + ggml_tensor * timestep_input_weight = nullptr; + ggml_tensor * timestep_input_bias = nullptr; ggml_tensor * timestep_output_weight = nullptr; - ggml_tensor * timestep_output_bias = nullptr; + ggml_tensor * timestep_output_bias = nullptr; std::vector blocks; - ggml_tensor * state_input_weight = nullptr; - ggml_tensor * state_input_bias = nullptr; - ggml_tensor * state_output_weight = nullptr; - ggml_tensor * state_output_bias = nullptr; - ggml_tensor * action_input_weight = nullptr; - ggml_tensor * action_input_bias = nullptr; + ggml_tensor * state_input_weight = nullptr; + ggml_tensor * state_input_bias = nullptr; + ggml_tensor * state_output_weight = nullptr; + ggml_tensor * state_output_bias = nullptr; + ggml_tensor * action_input_weight = nullptr; + ggml_tensor * action_input_bias = nullptr; ggml_tensor * action_time_mix_weight = nullptr; - ggml_tensor * action_time_mix_bias = nullptr; - ggml_tensor * action_output_weight = nullptr; - ggml_tensor * action_output_bias = nullptr; - ggml_tensor * velocity_input_weight = nullptr; - ggml_tensor * velocity_input_bias = nullptr; + ggml_tensor * action_time_mix_bias = nullptr; + ggml_tensor * action_output_weight = nullptr; + ggml_tensor * action_output_bias = nullptr; + ggml_tensor * velocity_input_weight = nullptr; + ggml_tensor * velocity_input_bias = nullptr; ggml_tensor * velocity_output_weight = nullptr; - ggml_tensor * velocity_output_bias = nullptr; - ggml_tensor * future_tokens = nullptr; - ggml_tensor * action_position = nullptr; + ggml_tensor * velocity_output_bias = nullptr; + ggml_tensor * future_tokens = nullptr; + ggml_tensor * action_position = nullptr; }; -int require_key(gguf_context * gguf, const char * key, gguf_type type) { - const int index = gguf_find_key(gguf, key); - if (index < 0) { - throw std::runtime_error(std::string("missing required StarVLA PI GGUF metadata: ") + - key); - } - if (gguf_get_kv_type(gguf, index) != type) { - throw std::runtime_error(std::string("invalid StarVLA PI GGUF metadata type: ") + - key); - } - return index; -} - -std::string require_string(gguf_context * gguf, const char * key) { - return gguf_get_val_str(gguf, require_key(gguf, key, GGUF_TYPE_STRING)); -} - -int require_i32(gguf_context * gguf, const char * key) { - return gguf_get_val_i32(gguf, require_key(gguf, key, GGUF_TYPE_INT32)); -} - -float require_f32(gguf_context * gguf, const char * key) { - return gguf_get_val_f32(gguf, require_key(gguf, key, GGUF_TYPE_FLOAT32)); -} - -bool require_bool(gguf_context * gguf, const char * key) { - return gguf_get_val_bool(gguf, require_key(gguf, key, GGUF_TYPE_BOOL)); -} - -int require_array(gguf_context * gguf, const char * key, gguf_type element_type) { - const int index = require_key(gguf, key, GGUF_TYPE_ARRAY); - if (gguf_get_arr_type(gguf, index) != element_type) { - throw std::runtime_error( - std::string("invalid StarVLA PI GGUF array element type: ") + key); - } - return index; -} +using detail::has_shape; +using detail::require_f32; +using detail::require_i32; +using detail::require_i32_array; +using detail::require_string; +using detail::require_string_array; -std::vector require_string_array(gguf_context * gguf, - const char * key) { - const int index = require_array(gguf, key, GGUF_TYPE_STRING); - const size_t count = gguf_get_arr_n(gguf, index); - std::vector result; - result.reserve(count); - for (size_t i = 0; i < count; ++i) { - result.emplace_back(gguf_get_arr_str(gguf, index, i)); - } - return result; -} - -std::vector require_i32_array(gguf_context * gguf, const char * key) { - const int index = require_array(gguf, key, GGUF_TYPE_INT32); - const size_t count = gguf_get_arr_n(gguf, index); - const auto * data = - static_cast(gguf_get_arr_data(gguf, index)); - if (data == nullptr && count != 0) { - throw std::runtime_error(std::string("missing StarVLA PI GGUF array data: ") + - key); - } - return count == 0 ? std::vector() - : std::vector(data, data + count); -} - -std::vector require_f32_array(gguf_context * gguf, const char * key) { - const int index = require_array(gguf, key, GGUF_TYPE_FLOAT32); - const size_t count = gguf_get_arr_n(gguf, index); - const auto * data = static_cast(gguf_get_arr_data(gguf, index)); - if (data == nullptr && count != 0) { - throw std::runtime_error(std::string("missing StarVLA PI GGUF array data: ") + - key); - } - return count == 0 ? std::vector() - : std::vector(data, data + count); -} - -std::vector require_bool_array(gguf_context * gguf, const char * key) { - const int index = require_array(gguf, key, GGUF_TYPE_BOOL); - const size_t count = gguf_get_arr_n(gguf, index); - const auto * data = - static_cast(gguf_get_arr_data(gguf, index)); - if (data == nullptr && count != 0) { - throw std::runtime_error(std::string("missing StarVLA PI GGUF array data: ") + - key); - } - std::vector result(count); - for (size_t i = 0; i < count; ++i) { - result[i] = data[i] != 0 ? 1 : 0; - } - return result; -} - -std::string profile_key(int profile_index, const char * suffix) { - return "starvla.normalization.profile." + std::to_string(profile_index) + "." + - suffix; -} - -bool has_shape(const ggml_tensor * tensor, - std::initializer_list expected) { - if (tensor == nullptr || - static_cast(ggml_n_dims(tensor)) != expected.size()) { - return false; - } - size_t dimension = 0; - for (const int64_t value : expected) { - if (tensor->ne[dimension++] != value) { - return false; - } - } - return true; -} - -const char * mode_name(backend_mode mode) { - switch (mode) { - case backend_mode::cpu: - return "cpu"; - case backend_mode::cuda: - return "cuda"; - case backend_mode::metal: - return "metal"; - } - return "unknown"; -} - -std::vector expected_hidden_tuple_indices(int qwen_layer_count, - int block_count) { +std::vector expected_hidden_tuple_indices(int qwen_layer_count, int block_count) { std::vector result; result.reserve(static_cast(block_count)); const int first = qwen_layer_count + 1 - block_count; @@ -198,187 +81,94 @@ std::vector expected_hidden_tuple_indices(int qwen_layer_count, class PIGGUFLoader final : public gguf_loader { public: - PIGGUFLoader(PIPolicyConfig & config, PIWeights & weights) - : config_(config), weights_(weights) {} + PIGGUFLoader(PIPolicyConfig & config, PIWeights & weights) : config_(config), weights_(weights) {} protected: bool parse_metadata(gguf_context * gguf) override { if (require_string(gguf, "general.architecture") != "starvla-policy" || - require_i32(gguf, "starvla.schema_version") != 1 || - require_string(gguf, "starvla.framework") != "pi") { + require_i32(gguf, "starvla.schema_version") != 1 || require_string(gguf, "starvla.framework") != "pi") { throw std::runtime_error("GGUF is not a supported StarVLA PI policy"); } - config_.backbone_arch = require_string(gguf, "starvla.backbone.arch"); - config_.bundle_uuid = require_string(gguf, "starvla.bundle.uuid"); - config_.text_filename = - require_string(gguf, "starvla.component.text.filename"); - config_.mmproj_filename = - require_string(gguf, "starvla.component.mmproj.filename"); - if (config_.backbone_arch != "qwen2_5_vl" || - config_.bundle_uuid.empty() || config_.text_filename.empty() || + config_.backbone_arch = require_string(gguf, "starvla.backbone.arch"); + config_.bundle_uuid = require_string(gguf, "starvla.bundle.uuid"); + config_.text_filename = require_string(gguf, "starvla.component.text.filename"); + config_.mmproj_filename = require_string(gguf, "starvla.component.mmproj.filename"); + if (config_.backbone_arch != "qwen2_5_vl" || config_.bundle_uuid.empty() || config_.text_filename.empty() || config_.mmproj_filename.empty()) { throw std::runtime_error("StarVLA PI bundle metadata is incomplete"); } - config_.qwen_hidden_dim = - require_i32(gguf, "starvla.qwen.hidden_size"); - config_.qwen_input_embedding_dim = - require_i32(gguf, "starvla.qwen.input_embedding_size"); - config_.qwen_layer_count = - require_i32(gguf, "starvla.qwen.layer_count"); - config_.qwen_vocab_size = require_i32(gguf, "starvla.qwen.vocab_size"); - config_.cot_template = - require_string(gguf, "starvla.prompt.cot_template"); - config_.qwen_hidden_tuple_indices = - require_i32_array(gguf, "starvla.conditioning.hidden_tuple_indices"); + config_.qwen_hidden_dim = require_i32(gguf, "starvla.qwen.hidden_size"); + config_.qwen_input_embedding_dim = require_i32(gguf, "starvla.qwen.input_embedding_size"); + config_.qwen_layer_count = require_i32(gguf, "starvla.qwen.layer_count"); + config_.qwen_vocab_size = require_i32(gguf, "starvla.qwen.vocab_size"); + config_.cot_template = require_string(gguf, "starvla.prompt.cot_template"); + config_.qwen_hidden_tuple_indices = require_i32_array(gguf, "starvla.conditioning.hidden_tuple_indices"); config_.image_count = require_i32(gguf, "starvla.image.count"); - config_.image_names = - require_string_array(gguf, "starvla.image.names"); + config_.image_names = require_string_array(gguf, "starvla.image.names"); config_.image_framework_inference_pre_resize_width = require_i32(gguf, "starvla.image.framework_inference_pre_resize_width"); config_.image_framework_inference_pre_resize_height = require_i32(gguf, "starvla.image.framework_inference_pre_resize_height"); - config_.image_processor_min_pixels = - require_i32(gguf, "starvla.image.processor_min_pixels"); - config_.image_processor_max_pixels = - require_i32(gguf, "starvla.image.processor_max_pixels"); - config_.image_patch_size = - require_i32(gguf, "starvla.image.patch_size"); - config_.image_spatial_merge_size = - require_i32(gguf, "starvla.image.spatial_merge_size"); - config_.image_min_token_count = - require_i32(gguf, "starvla.image.min_token_count"); - config_.image_max_token_count = - require_i32(gguf, "starvla.image.max_token_count"); - - config_.dit_width = require_i32(gguf, "starvla.pi.dit_width"); - config_.block_count = require_i32(gguf, "starvla.pi.block_count"); - config_.attention_head_count = - require_i32(gguf, "starvla.pi.attention_head_count"); - config_.attention_head_dim = - require_i32(gguf, "starvla.pi.attention_head_dim"); - config_.cross_attention_dim = - require_i32(gguf, "starvla.pi.cross_attention_dim"); - config_.feed_forward_dim = - require_i32(gguf, "starvla.pi.feed_forward_dim"); - config_.mlp_hidden_dim = - require_i32(gguf, "starvla.pi.mlp_hidden_dimension"); - config_.state_dim = require_i32(gguf, "starvla.state.dimension"); - config_.action_dim = require_i32(gguf, "starvla.action.dimension"); - config_.horizon = require_i32(gguf, "starvla.action.horizon"); - config_.state_token_count = - require_i32(gguf, "starvla.pi.state_token_count"); - config_.future_token_count = - require_i32(gguf, "starvla.pi.future_token_count"); - config_.action_position_count = - require_i32(gguf, "starvla.pi.action_position_count"); - config_.timestep_projection_dim = - require_i32(gguf, "starvla.pi.timestep_projection_dim"); - config_.num_inference_timesteps = - require_i32(gguf, "starvla.pi.num_inference_timesteps"); - config_.ada_norm_epsilon = - require_f32(gguf, "starvla.pi.ada_norm_epsilon"); - config_.euler_dt = require_f32(gguf, "starvla.pi.euler_dt"); - config_.timestep_ids = - require_i32_array(gguf, "starvla.pi.timestep_ids"); + config_.image_processor_min_pixels = require_i32(gguf, "starvla.image.processor_min_pixels"); + config_.image_processor_max_pixels = require_i32(gguf, "starvla.image.processor_max_pixels"); + config_.image_patch_size = require_i32(gguf, "starvla.image.patch_size"); + config_.image_spatial_merge_size = require_i32(gguf, "starvla.image.spatial_merge_size"); + config_.image_min_token_count = require_i32(gguf, "starvla.image.min_token_count"); + config_.image_max_token_count = require_i32(gguf, "starvla.image.max_token_count"); + + config_.dit_width = require_i32(gguf, "starvla.pi.dit_width"); + config_.block_count = require_i32(gguf, "starvla.pi.block_count"); + config_.attention_head_count = require_i32(gguf, "starvla.pi.attention_head_count"); + config_.attention_head_dim = require_i32(gguf, "starvla.pi.attention_head_dim"); + config_.cross_attention_dim = require_i32(gguf, "starvla.pi.cross_attention_dim"); + config_.feed_forward_dim = require_i32(gguf, "starvla.pi.feed_forward_dim"); + config_.mlp_hidden_dim = require_i32(gguf, "starvla.pi.mlp_hidden_dimension"); + config_.state_dim = require_i32(gguf, "starvla.state.dimension"); + config_.action_dim = require_i32(gguf, "starvla.action.dimension"); + config_.horizon = require_i32(gguf, "starvla.action.horizon"); + config_.state_token_count = require_i32(gguf, "starvla.pi.state_token_count"); + config_.future_token_count = require_i32(gguf, "starvla.pi.future_token_count"); + config_.action_position_count = require_i32(gguf, "starvla.pi.action_position_count"); + config_.timestep_projection_dim = require_i32(gguf, "starvla.pi.timestep_projection_dim"); + config_.num_inference_timesteps = require_i32(gguf, "starvla.pi.num_inference_timesteps"); + config_.ada_norm_epsilon = require_f32(gguf, "starvla.pi.ada_norm_epsilon"); + config_.euler_dt = require_f32(gguf, "starvla.pi.euler_dt"); + config_.timestep_ids = require_i32_array(gguf, "starvla.pi.timestep_ids"); const std::vector expected_indices = - expected_hidden_tuple_indices(config_.qwen_layer_count, - config_.block_count); + expected_hidden_tuple_indices(config_.qwen_layer_count, config_.block_count); const bool dimensions_valid = - config_.qwen_hidden_dim > 0 && - config_.qwen_input_embedding_dim == config_.qwen_hidden_dim && - config_.qwen_layer_count >= config_.block_count && - config_.qwen_vocab_size > 0 && config_.dit_width > 0 && - config_.dit_width % 2 == 0 && config_.block_count > 0 && - config_.attention_head_count > 0 && + config_.qwen_hidden_dim > 0 && config_.qwen_input_embedding_dim == config_.qwen_hidden_dim && + config_.qwen_layer_count >= config_.block_count && config_.qwen_vocab_size > 0 && config_.dit_width > 0 && + config_.dit_width % 2 == 0 && config_.block_count > 0 && config_.attention_head_count > 0 && config_.attention_head_dim > 0 && - config_.attention_head_count * config_.attention_head_dim == - config_.dit_width && - config_.cross_attention_dim == config_.qwen_hidden_dim && - config_.feed_forward_dim > 0 && - config_.mlp_hidden_dim > 0 && config_.state_dim > 0 && - config_.action_dim > 0 && config_.horizon > 0 && - config_.state_token_count == 1 && - config_.future_token_count > 0 && - config_.action_position_count >= config_.horizon && - config_.timestep_projection_dim >= 4 && - config_.timestep_projection_dim % 2 == 0 && - config_.num_inference_timesteps > 0 && - config_.timestep_ids.size() == - static_cast(config_.num_inference_timesteps) && - std::isfinite(config_.ada_norm_epsilon) && - config_.ada_norm_epsilon > 0.0f && std::isfinite(config_.euler_dt) && - config_.euler_dt > 0.0f && - config_.qwen_hidden_tuple_indices == expected_indices && - config_.image_count > 0 && + config_.attention_head_count * config_.attention_head_dim == config_.dit_width && + config_.cross_attention_dim == config_.qwen_hidden_dim && config_.feed_forward_dim > 0 && + config_.mlp_hidden_dim > 0 && config_.state_dim > 0 && config_.action_dim > 0 && config_.horizon > 0 && + config_.state_token_count == 1 && config_.future_token_count > 0 && + config_.action_position_count >= config_.horizon && config_.timestep_projection_dim >= 4 && + config_.timestep_projection_dim % 2 == 0 && config_.num_inference_timesteps > 0 && + config_.timestep_ids.size() == static_cast(config_.num_inference_timesteps) && + std::isfinite(config_.ada_norm_epsilon) && config_.ada_norm_epsilon > 0.0f && + std::isfinite(config_.euler_dt) && config_.euler_dt > 0.0f && + config_.qwen_hidden_tuple_indices == expected_indices && config_.image_count > 0 && config_.image_names.size() == static_cast(config_.image_count) && config_.image_framework_inference_pre_resize_width > 0 && - config_.image_framework_inference_pre_resize_height > 0 && - config_.image_processor_min_pixels > 0 && - config_.image_processor_max_pixels >= - config_.image_processor_min_pixels && - config_.image_patch_size > 0 && - config_.image_spatial_merge_size > 0 && - config_.image_min_token_count > 0 && - config_.image_max_token_count >= - config_.image_min_token_count && - !config_.cot_template.empty(); + config_.image_framework_inference_pre_resize_height > 0 && config_.image_processor_min_pixels > 0 && + config_.image_processor_max_pixels >= config_.image_processor_min_pixels && config_.image_patch_size > 0 && + config_.image_spatial_merge_size > 0 && config_.image_min_token_count > 0 && + config_.image_max_token_count >= config_.image_min_token_count && !config_.cot_template.empty(); if (!dimensions_valid) { - throw std::runtime_error( - "StarVLA PI dimensions, hidden taps, or sampler schedule are incompatible"); + throw std::runtime_error("StarVLA PI dimensions, hidden taps, or sampler schedule are incompatible"); } NormalizationConfig & normalization = config_.normalization; - normalization.clip_actions = - require_bool(gguf, "starvla.normalization.clip_actions"); - normalization.binary_threshold = - require_f32(gguf, "starvla.normalization.binary_threshold"); - normalization.binary_comparison = - require_string(gguf, "starvla.normalization.binary_comparison"); - normalization.continuous_dimensions = - require_i32_array(gguf, "starvla.action.continuous_dimensions"); - normalization.binary_dimensions = - require_i32_array(gguf, "starvla.action.binary_dimensions"); - const int profile_count = - require_i32(gguf, "starvla.normalization.profile_count"); - const std::vector keys = - require_string_array(gguf, "starvla.normalization.profile_keys"); - if (profile_count <= 0 || - keys.size() != static_cast(profile_count)) { - throw std::runtime_error( - "StarVLA PI normalization profile count is inconsistent"); - } - normalization.profiles.clear(); - normalization.profiles.reserve(static_cast(profile_count)); - for (int i = 0; i < profile_count; ++i) { - NormalizationProfile profile; - profile.key = - require_string(gguf, profile_key(i, "key").c_str()); - profile.action_q01 = - require_f32_array(gguf, profile_key(i, "action_q01").c_str()); - profile.action_q99 = - require_f32_array(gguf, profile_key(i, "action_q99").c_str()); - profile.action_mask = - require_bool_array(gguf, - profile_key(i, "action_mask").c_str()); - if (profile.key != keys[static_cast(i)]) { - throw std::runtime_error( - "StarVLA PI normalization profile order is inconsistent"); - } - normalization.profiles.push_back(std::move(profile)); - } - std::string normalization_error; - if (!validate_normalization_config(normalization, config_.action_dim, - normalization_error)) { - throw std::runtime_error(normalization_error); - } - if (!normalization.clip_actions || - normalization.binary_comparison != "ge") { - throw std::runtime_error( - "StarVLA PI normalization must clip actions and use " - "binary comparison 'ge'"); + normalization = detail::require_normalization(gguf, config_.action_dim); + if (!normalization.clip_actions || normalization.binary_comparison != "ge") { + throw std::runtime_error("StarVLA PI normalization must clip actions and use " + "binary comparison 'ge'"); } return true; } @@ -387,19 +177,14 @@ class PIGGUFLoader final : public gguf_loader { auto bind = [&](ggml_tensor *& destination, const std::string & name) { destination = require_tensor(ctx_data, name); }; - bind(weights_.timestep_input_weight, - "starvla.policy.pi.timestep.input.weight"); - bind(weights_.timestep_input_bias, - "starvla.policy.pi.timestep.input.bias"); - bind(weights_.timestep_output_weight, - "starvla.policy.pi.timestep.output.weight"); - bind(weights_.timestep_output_bias, - "starvla.policy.pi.timestep.output.bias"); + bind(weights_.timestep_input_weight, "starvla.policy.pi.timestep.input.weight"); + bind(weights_.timestep_input_bias, "starvla.policy.pi.timestep.input.bias"); + bind(weights_.timestep_output_weight, "starvla.policy.pi.timestep.output.weight"); + bind(weights_.timestep_output_bias, "starvla.policy.pi.timestep.output.bias"); weights_.blocks.clear(); weights_.blocks.reserve(static_cast(config_.block_count)); for (int block = 0; block < config_.block_count; ++block) { - const std::string prefix = - "starvla.policy.pi.block." + std::to_string(block) + "."; + const std::string prefix = "starvla.policy.pi.block." + std::to_string(block) + "."; PIBlockWeights current; bind(current.ada_norm_weight, prefix + "ada_norm.weight"); bind(current.ada_norm_bias, prefix + "ada_norm.bias"); @@ -409,105 +194,67 @@ class PIGGUFLoader final : public gguf_loader { bind(current.key_bias, prefix + "attention.key.bias"); bind(current.value_weight, prefix + "attention.value.weight"); bind(current.value_bias, prefix + "attention.value.bias"); - bind(current.attention_output_weight, - prefix + "attention.output.weight"); - bind(current.attention_output_bias, - prefix + "attention.output.bias"); - bind(current.feed_forward_input_weight, - prefix + "feed_forward.input.weight"); - bind(current.feed_forward_input_bias, - prefix + "feed_forward.input.bias"); - bind(current.feed_forward_output_weight, - prefix + "feed_forward.output.weight"); - bind(current.feed_forward_output_bias, - prefix + "feed_forward.output.bias"); + bind(current.attention_output_weight, prefix + "attention.output.weight"); + bind(current.attention_output_bias, prefix + "attention.output.bias"); + bind(current.feed_forward_input_weight, prefix + "feed_forward.input.weight"); + bind(current.feed_forward_input_bias, prefix + "feed_forward.input.bias"); + bind(current.feed_forward_output_weight, prefix + "feed_forward.output.weight"); + bind(current.feed_forward_output_bias, prefix + "feed_forward.output.bias"); weights_.blocks.push_back(current); } - bind(weights_.state_input_weight, - "starvla.policy.pi.state.input.weight"); + bind(weights_.state_input_weight, "starvla.policy.pi.state.input.weight"); bind(weights_.state_input_bias, "starvla.policy.pi.state.input.bias"); - bind(weights_.state_output_weight, - "starvla.policy.pi.state.output.weight"); + bind(weights_.state_output_weight, "starvla.policy.pi.state.output.weight"); bind(weights_.state_output_bias, "starvla.policy.pi.state.output.bias"); - bind(weights_.action_input_weight, - "starvla.policy.pi.action.input.weight"); + bind(weights_.action_input_weight, "starvla.policy.pi.action.input.weight"); bind(weights_.action_input_bias, "starvla.policy.pi.action.input.bias"); - bind(weights_.action_time_mix_weight, - "starvla.policy.pi.action.time_mix.weight"); - bind(weights_.action_time_mix_bias, - "starvla.policy.pi.action.time_mix.bias"); - bind(weights_.action_output_weight, - "starvla.policy.pi.action.output.weight"); - bind(weights_.action_output_bias, - "starvla.policy.pi.action.output.bias"); - bind(weights_.velocity_input_weight, - "starvla.policy.pi.velocity.input.weight"); - bind(weights_.velocity_input_bias, - "starvla.policy.pi.velocity.input.bias"); - bind(weights_.velocity_output_weight, - "starvla.policy.pi.velocity.output.weight"); - bind(weights_.velocity_output_bias, - "starvla.policy.pi.velocity.output.bias"); + bind(weights_.action_time_mix_weight, "starvla.policy.pi.action.time_mix.weight"); + bind(weights_.action_time_mix_bias, "starvla.policy.pi.action.time_mix.bias"); + bind(weights_.action_output_weight, "starvla.policy.pi.action.output.weight"); + bind(weights_.action_output_bias, "starvla.policy.pi.action.output.bias"); + bind(weights_.velocity_input_weight, "starvla.policy.pi.velocity.input.weight"); + bind(weights_.velocity_input_bias, "starvla.policy.pi.velocity.input.bias"); + bind(weights_.velocity_output_weight, "starvla.policy.pi.velocity.output.weight"); + bind(weights_.velocity_output_bias, "starvla.policy.pi.velocity.output.bias"); bind(weights_.future_tokens, "starvla.policy.pi.future_tokens.weight"); - bind(weights_.action_position, - "starvla.policy.pi.action_position.weight"); + bind(weights_.action_position, "starvla.policy.pi.action_position.weight"); const int64_t width = config_.dit_width; - if (!has_shape(weights_.timestep_input_weight, - {config_.timestep_projection_dim, width}) || + if (!has_shape(weights_.timestep_input_weight, {config_.timestep_projection_dim, width}) || !has_shape(weights_.timestep_input_bias, {width}) || !has_shape(weights_.timestep_output_weight, {width, width}) || !has_shape(weights_.timestep_output_bias, {width}) || - !has_shape(weights_.state_input_weight, - {config_.state_dim, config_.mlp_hidden_dim}) || + !has_shape(weights_.state_input_weight, {config_.state_dim, config_.mlp_hidden_dim}) || !has_shape(weights_.state_input_bias, {config_.mlp_hidden_dim}) || - !has_shape(weights_.state_output_weight, - {config_.mlp_hidden_dim, width}) || + !has_shape(weights_.state_output_weight, {config_.mlp_hidden_dim, width}) || !has_shape(weights_.state_output_bias, {width}) || - !has_shape(weights_.action_input_weight, - {config_.action_dim, width}) || + !has_shape(weights_.action_input_weight, {config_.action_dim, width}) || !has_shape(weights_.action_input_bias, {width}) || - !has_shape(weights_.action_time_mix_weight, - {2 * width, width}) || + !has_shape(weights_.action_time_mix_weight, {2 * width, width}) || !has_shape(weights_.action_time_mix_bias, {width}) || !has_shape(weights_.action_output_weight, {width, width}) || !has_shape(weights_.action_output_bias, {width}) || - !has_shape(weights_.velocity_input_weight, - {width, config_.mlp_hidden_dim}) || - !has_shape(weights_.velocity_input_bias, - {config_.mlp_hidden_dim}) || - !has_shape(weights_.velocity_output_weight, - {config_.mlp_hidden_dim, config_.action_dim}) || + !has_shape(weights_.velocity_input_weight, {width, config_.mlp_hidden_dim}) || + !has_shape(weights_.velocity_input_bias, {config_.mlp_hidden_dim}) || + !has_shape(weights_.velocity_output_weight, {config_.mlp_hidden_dim, config_.action_dim}) || !has_shape(weights_.velocity_output_bias, {config_.action_dim}) || - !has_shape(weights_.future_tokens, - {width, config_.future_token_count}) || - !has_shape(weights_.action_position, - {width, config_.action_position_count})) { - throw std::runtime_error( - "StarVLA PI non-transformer tensor has an incompatible ggml shape"); + !has_shape(weights_.future_tokens, {width, config_.future_token_count}) || + !has_shape(weights_.action_position, {width, config_.action_position_count})) { + throw std::runtime_error("StarVLA PI non-transformer tensor has an incompatible ggml shape"); } for (const PIBlockWeights & block : weights_.blocks) { - if (!has_shape(block.ada_norm_weight, {width, 2 * width}) || - !has_shape(block.ada_norm_bias, {2 * width}) || - !has_shape(block.query_weight, {width, width}) || - !has_shape(block.query_bias, {width}) || - !has_shape(block.key_weight, - {config_.cross_attention_dim, width}) || + if (!has_shape(block.ada_norm_weight, {width, 2 * width}) || !has_shape(block.ada_norm_bias, {2 * width}) || + !has_shape(block.query_weight, {width, width}) || !has_shape(block.query_bias, {width}) || + !has_shape(block.key_weight, {config_.cross_attention_dim, width}) || !has_shape(block.key_bias, {width}) || - !has_shape(block.value_weight, - {config_.cross_attention_dim, width}) || - !has_shape(block.value_bias, {width}) || - !has_shape(block.attention_output_weight, {width, width}) || + !has_shape(block.value_weight, {config_.cross_attention_dim, width}) || + !has_shape(block.value_bias, {width}) || !has_shape(block.attention_output_weight, {width, width}) || !has_shape(block.attention_output_bias, {width}) || - !has_shape(block.feed_forward_input_weight, - {width, config_.feed_forward_dim}) || - !has_shape(block.feed_forward_input_bias, - {config_.feed_forward_dim}) || - !has_shape(block.feed_forward_output_weight, - {config_.feed_forward_dim, width}) || + !has_shape(block.feed_forward_input_weight, {width, config_.feed_forward_dim}) || + !has_shape(block.feed_forward_input_bias, {config_.feed_forward_dim}) || + !has_shape(block.feed_forward_output_weight, {config_.feed_forward_dim, width}) || !has_shape(block.feed_forward_output_bias, {width})) { - throw std::runtime_error( - "StarVLA PI transformer tensor has an incompatible ggml shape"); + throw std::runtime_error("StarVLA PI transformer tensor has an incompatible ggml shape"); } } return true; @@ -519,20 +266,15 @@ class PIGGUFLoader final : public gguf_loader { }; std::vector timestep_projection_table(const PIPolicyConfig & config) { - std::vector result( - static_cast(config.num_inference_timesteps) * - config.timestep_projection_dim); + std::vector result(static_cast(config.num_inference_timesteps) * config.timestep_projection_dim); const int half = config.timestep_projection_dim / 2; for (int step = 0; step < config.num_inference_timesteps; ++step) { - const float timestep = - static_cast(config.timestep_ids[static_cast(step)]); + const float timestep = static_cast(config.timestep_ids[static_cast(step)]); for (int i = 0; i < half; ++i) { - const float exponent = - -std::log(10000.0f) * i / static_cast(half - 1); - const float angle = timestep * std::exp(exponent); - const size_t offset = - static_cast(step) * config.timestep_projection_dim; - result[offset + static_cast(i)] = std::cos(angle); + const float exponent = -std::log(10000.0f) * i / static_cast(half - 1); + const float angle = timestep * std::exp(exponent); + const size_t offset = static_cast(step) * config.timestep_projection_dim; + result[offset + static_cast(i)] = std::cos(angle); result[offset + static_cast(i + half)] = std::sin(angle); } } @@ -540,19 +282,15 @@ std::vector timestep_projection_table(const PIPolicyConfig & config) { } std::vector action_time_table(const PIPolicyConfig & config) { - std::vector result( - static_cast(config.num_inference_timesteps) * config.dit_width); + std::vector result(static_cast(config.num_inference_timesteps) * config.dit_width); const int half = config.dit_width / 2; for (int step = 0; step < config.num_inference_timesteps; ++step) { - const float timestep = - static_cast(config.timestep_ids[static_cast(step)]); + const float timestep = static_cast(config.timestep_ids[static_cast(step)]); for (int i = 0; i < half; ++i) { - const float exponent = - -std::log(10000.0f) * i / static_cast(half); - const float angle = timestep * std::exp(exponent); - const size_t offset = - static_cast(step) * config.dit_width; - result[offset + static_cast(i)] = std::sin(angle); + const float exponent = -std::log(10000.0f) * i / static_cast(half); + const float angle = timestep * std::exp(exponent); + const size_t offset = static_cast(step) * config.dit_width; + result[offset + static_cast(i)] = std::sin(angle); result[offset + static_cast(i + half)] = std::cos(angle); } } @@ -569,21 +307,20 @@ struct PIPolicy::Impl { std::vector backends; ggml_backend_sched_t scheduler = nullptr; backend_buft_policy buft_policy; - backend_mode mode = backend_mode::cpu; - int n_threads = 0; - int verbosity = 0; - ggml_context * graph_context = nullptr; - ggml_cgraph * graph = nullptr; - ggml_tensor * hidden_input = nullptr; - ggml_tensor * state_input = nullptr; - ggml_tensor * noise_input = nullptr; + backend_mode mode = backend_mode::cpu; + int n_threads = 0; + int verbosity = 0; + ggml_context * graph_context = nullptr; + ggml_cgraph * graph = nullptr; + ggml_tensor * hidden_input = nullptr; + ggml_tensor * state_input = nullptr; + ggml_tensor * noise_input = nullptr; ggml_tensor * timestep_projection_input = nullptr; - ggml_tensor * action_time_input = nullptr; - ggml_tensor * scalar_one_input = nullptr; - ggml_tensor * output = nullptr; - size_t conditioning_token_count = 0; - bool graph_uses_state = false; - size_t graph_builds = 0; + ggml_tensor * action_time_input = nullptr; + ggml_tensor * scalar_one_input = nullptr; + ggml_tensor * output = nullptr; + size_t conditioning_token_count = 0; + bool graph_uses_state = false; std::vector timestep_table; std::vector action_table; @@ -624,72 +361,56 @@ struct PIPolicy::Impl { ggml_free(graph_context); graph_context = nullptr; } - graph = nullptr; - hidden_input = nullptr; - state_input = nullptr; - noise_input = nullptr; + graph = nullptr; + hidden_input = nullptr; + state_input = nullptr; + noise_input = nullptr; timestep_projection_input = nullptr; - action_time_input = nullptr; - scalar_one_input = nullptr; - output = nullptr; - conditioning_token_count = 0; - graph_uses_state = false; + action_time_input = nullptr; + scalar_one_input = nullptr; + output = nullptr; + conditioning_token_count = 0; + graph_uses_state = false; } void build_graph(size_t token_count, bool include_state) { clear_graph(); - if (token_count == 0 || - token_count > static_cast(std::numeric_limits::max())) { - throw std::runtime_error( - "invalid StarVLA PI conditioning token count"); + if (token_count == 0 || token_count > static_cast(std::numeric_limits::max())) { + throw std::runtime_error("invalid StarVLA PI conditioning token count"); } ggml_init_params params{}; - params.mem_size = - kGraphSize * ggml_tensor_overhead() + - ggml_graph_overhead_custom(kGraphSize, false); + params.mem_size = kGraphSize * ggml_tensor_overhead() + ggml_graph_overhead_custom(kGraphSize, false); params.mem_buffer = nullptr; - params.no_alloc = true; - graph_context = ggml_init(params); + params.no_alloc = true; + graph_context = ggml_init(params); if (graph_context == nullptr) { - throw std::runtime_error( - "failed to initialize StarVLA PI graph context"); + throw std::runtime_error("failed to initialize StarVLA PI graph context"); } - const int width = config.dit_width; - const int heads = config.attention_head_count; + const int width = config.dit_width; + const int heads = config.attention_head_count; const int head_dim = config.attention_head_dim; - hidden_input = ggml_new_tensor_3d( - graph_context, GGML_TYPE_F32, config.qwen_hidden_dim, - static_cast(token_count), config.block_count); + hidden_input = ggml_new_tensor_3d(graph_context, GGML_TYPE_F32, config.qwen_hidden_dim, + static_cast(token_count), config.block_count); if (include_state) { - state_input = ggml_new_tensor_1d(graph_context, GGML_TYPE_F32, - config.state_dim); + state_input = ggml_new_tensor_1d(graph_context, GGML_TYPE_F32, config.state_dim); } - noise_input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, - config.action_dim, config.horizon); - timestep_projection_input = ggml_new_tensor_2d( - graph_context, GGML_TYPE_F32, config.timestep_projection_dim, - config.num_inference_timesteps); - action_time_input = ggml_new_tensor_2d( - graph_context, GGML_TYPE_F32, width, - config.num_inference_timesteps); - scalar_one_input = - ggml_new_tensor_1d(graph_context, GGML_TYPE_F32, 1); - if (hidden_input == nullptr || - (include_state && state_input == nullptr) || - noise_input == nullptr || timestep_projection_input == nullptr || - action_time_input == nullptr || scalar_one_input == nullptr) { - throw std::runtime_error( - "failed to create StarVLA PI graph inputs"); + noise_input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, config.action_dim, config.horizon); + timestep_projection_input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, config.timestep_projection_dim, + config.num_inference_timesteps); + action_time_input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, width, config.num_inference_timesteps); + scalar_one_input = ggml_new_tensor_1d(graph_context, GGML_TYPE_F32, 1); + if (hidden_input == nullptr || (include_state && state_input == nullptr) || noise_input == nullptr || + timestep_projection_input == nullptr || action_time_input == nullptr || scalar_one_input == nullptr) { + throw std::runtime_error("failed to create StarVLA PI graph inputs"); } ggml_set_name(hidden_input, "starvla_pi_qwen_hidden_states"); if (state_input != nullptr) { ggml_set_name(state_input, "starvla_pi_state"); } ggml_set_name(noise_input, "starvla_pi_initial_noise"); - ggml_set_name(timestep_projection_input, - "starvla_pi_timestep_projection_table"); + ggml_set_name(timestep_projection_input, "starvla_pi_timestep_projection_table"); ggml_set_name(action_time_input, "starvla_pi_action_time_table"); ggml_set_name(scalar_one_input, "starvla_pi_scalar_one"); ggml_set_input(hidden_input); @@ -702,170 +423,106 @@ struct PIPolicy::Impl { ggml_set_input(scalar_one_input); auto f32 = [&](ggml_tensor * tensor) { - return tensor->type == GGML_TYPE_F32 - ? tensor - : ggml_cast(graph_context, tensor, GGML_TYPE_F32); + return tensor->type == GGML_TYPE_F32 ? tensor : ggml_cast(graph_context, tensor, GGML_TYPE_F32); }; - auto linear = [&](ggml_tensor * value, ggml_tensor * weight, - ggml_tensor * bias) { - ggml_tensor * projected = - ggml_mul_mat(graph_context, weight, value); + auto linear = [&](ggml_tensor * value, ggml_tensor * weight, ggml_tensor * bias) { + ggml_tensor * projected = ggml_mul_mat(graph_context, weight, value); ggml_mul_mat_set_prec(projected, GGML_PREC_F32); return ggml_add(graph_context, projected, f32(bias)); }; - auto ada_norm = [&](ggml_tensor * value, ggml_tensor * temb, - const PIBlockWeights & block) { + auto ada_norm = [&](ggml_tensor * value, ggml_tensor * temb, const PIBlockWeights & block) { ggml_tensor * modulation = - linear(ggml_silu(graph_context, temb), - block.ada_norm_weight, block.ada_norm_bias); - ggml_tensor * scale = - ggml_view_1d(graph_context, modulation, width, 0); - ggml_tensor * shift = ggml_view_1d( - graph_context, modulation, width, - static_cast(width) * sizeof(float)); - ggml_tensor * normalized = - ggml_norm(graph_context, value, config.ada_norm_epsilon); - return ggml_add( - graph_context, - ggml_mul(graph_context, normalized, - ggml_add(graph_context, scale, scalar_one_input)), - shift); + linear(ggml_silu(graph_context, temb), block.ada_norm_weight, block.ada_norm_bias); + ggml_tensor * scale = ggml_view_1d(graph_context, modulation, width, 0); + ggml_tensor * shift = + ggml_view_1d(graph_context, modulation, width, static_cast(width) * sizeof(float)); + ggml_tensor * normalized = ggml_norm(graph_context, value, config.ada_norm_epsilon); + return ggml_add(graph_context, + ggml_mul(graph_context, normalized, ggml_add(graph_context, scale, scalar_one_input)), + shift); }; - auto attention = [&](ggml_tensor * query_source, - ggml_tensor * key_value_source, - const PIBlockWeights & block) { - const int64_t query_count = query_source->ne[1]; + auto attention = [&](ggml_tensor * query_source, ggml_tensor * key_value_source, const PIBlockWeights & block) { + const int64_t query_count = query_source->ne[1]; const int64_t key_value_count = key_value_source->ne[1]; - ggml_tensor * query = - linear(query_source, block.query_weight, block.query_bias); - ggml_tensor * key = - linear(key_value_source, block.key_weight, block.key_bias); - ggml_tensor * value = - linear(key_value_source, block.value_weight, block.value_bias); - query = ggml_reshape_3d(graph_context, query, head_dim, heads, - query_count); - key = ggml_reshape_3d(graph_context, key, head_dim, heads, - key_value_count); - value = ggml_reshape_3d(graph_context, value, head_dim, heads, - key_value_count); - query = ggml_permute(graph_context, query, 0, 2, 1, 3); - key = ggml_permute(graph_context, key, 0, 2, 1, 3); - value = ggml_cont( - graph_context, - ggml_permute(graph_context, value, 1, 2, 0, 3)); - ggml_tensor * scores = ggml_mul_mat(graph_context, key, query); + ggml_tensor * query = linear(query_source, block.query_weight, block.query_bias); + ggml_tensor * key = linear(key_value_source, block.key_weight, block.key_bias); + ggml_tensor * value = linear(key_value_source, block.value_weight, block.value_bias); + query = ggml_reshape_3d(graph_context, query, head_dim, heads, query_count); + key = ggml_reshape_3d(graph_context, key, head_dim, heads, key_value_count); + value = ggml_reshape_3d(graph_context, value, head_dim, heads, key_value_count); + query = ggml_permute(graph_context, query, 0, 2, 1, 3); + key = ggml_permute(graph_context, key, 0, 2, 1, 3); + value = ggml_cont(graph_context, ggml_permute(graph_context, value, 1, 2, 0, 3)); + ggml_tensor * scores = ggml_mul_mat(graph_context, key, query); ggml_mul_mat_set_prec(scores, GGML_PREC_F32); - scores = ggml_soft_max_ext( - graph_context, scores, nullptr, - 1.0f / std::sqrt(static_cast(head_dim)), 0.0f); - ggml_tensor * attended = - ggml_mul_mat(graph_context, value, scores); + scores = + ggml_soft_max_ext(graph_context, scores, nullptr, 1.0f / std::sqrt(static_cast(head_dim)), 0.0f); + ggml_tensor * attended = ggml_mul_mat(graph_context, value, scores); ggml_mul_mat_set_prec(attended, GGML_PREC_F32); - attended = - ggml_permute(graph_context, attended, 0, 2, 1, 3); - attended = - ggml_cont_2d(graph_context, attended, width, query_count); - return linear(attended, block.attention_output_weight, - block.attention_output_bias); + attended = ggml_permute(graph_context, attended, 0, 2, 1, 3); + attended = ggml_cont_2d(graph_context, attended, width, query_count); + return linear(attended, block.attention_output_weight, block.attention_output_bias); }; ggml_tensor * state_features = nullptr; if (include_state) { state_features = - ggml_relu(graph_context, - linear(state_input, weights.state_input_weight, - weights.state_input_bias)); - state_features = - linear(state_features, weights.state_output_weight, - weights.state_output_bias); - state_features = - ggml_reshape_2d(graph_context, state_features, width, 1); + ggml_relu(graph_context, linear(state_input, weights.state_input_weight, weights.state_input_bias)); + state_features = linear(state_features, weights.state_output_weight, weights.state_output_bias); + state_features = ggml_reshape_2d(graph_context, state_features, width, 1); } - ggml_tensor * future = f32(weights.future_tokens); - ggml_tensor * position_view = ggml_view_2d( - graph_context, weights.action_position, width, config.horizon, - weights.action_position->nb[1], 0); - ggml_tensor * position = f32(position_view); - ggml_tensor * actions = noise_input; + ggml_tensor * future = f32(weights.future_tokens); + ggml_tensor * position_view = ggml_view_2d(graph_context, weights.action_position, width, config.horizon, + weights.action_position->nb[1], 0); + ggml_tensor * position = f32(position_view); + ggml_tensor * actions = noise_input; for (int step = 0; step < config.num_inference_timesteps; ++step) { - ggml_tensor * timestep_projection = ggml_view_1d( - graph_context, timestep_projection_input, - config.timestep_projection_dim, - static_cast(step) * - config.timestep_projection_dim * sizeof(float)); + ggml_tensor * timestep_projection = + ggml_view_1d(graph_context, timestep_projection_input, config.timestep_projection_dim, + static_cast(step) * config.timestep_projection_dim * sizeof(float)); ggml_tensor * temb = - linear(timestep_projection, weights.timestep_input_weight, - weights.timestep_input_bias); + linear(timestep_projection, weights.timestep_input_weight, weights.timestep_input_bias); temb = ggml_silu(graph_context, temb); - temb = linear(temb, weights.timestep_output_weight, - weights.timestep_output_bias); - - ggml_tensor * action_features = - linear(actions, weights.action_input_weight, - weights.action_input_bias); - ggml_tensor * action_time = ggml_view_1d( - graph_context, action_time_input, width, - static_cast(step) * width * sizeof(float)); - action_time = - ggml_repeat(graph_context, action_time, action_features); - action_features = - ggml_concat(graph_context, action_features, action_time, 0); - action_features = - linear(action_features, weights.action_time_mix_weight, - weights.action_time_mix_bias); + temb = linear(temb, weights.timestep_output_weight, weights.timestep_output_bias); + + ggml_tensor * action_features = linear(actions, weights.action_input_weight, weights.action_input_bias); + ggml_tensor * action_time = ggml_view_1d(graph_context, action_time_input, width, + static_cast(step) * width * sizeof(float)); + action_time = ggml_repeat(graph_context, action_time, action_features); + action_features = ggml_concat(graph_context, action_features, action_time, 0); + action_features = linear(action_features, weights.action_time_mix_weight, weights.action_time_mix_bias); action_features = ggml_silu(graph_context, action_features); - action_features = - linear(action_features, weights.action_output_weight, - weights.action_output_bias); - action_features = - ggml_add(graph_context, action_features, position); + action_features = linear(action_features, weights.action_output_weight, weights.action_output_bias); + action_features = ggml_add(graph_context, action_features, position); ggml_tensor * hidden = future; if (state_features != nullptr) { - hidden = - ggml_concat(graph_context, state_features, hidden, 1); + hidden = ggml_concat(graph_context, state_features, hidden, 1); } hidden = ggml_concat(graph_context, hidden, action_features, 1); - for (int block_index = 0; block_index < config.block_count; - ++block_index) { - const PIBlockWeights & block = - weights.blocks[static_cast(block_index)]; - ggml_tensor * layer_hidden = ggml_view_2d( - graph_context, hidden_input, config.qwen_hidden_dim, - static_cast(token_count), hidden_input->nb[1], - static_cast(block_index) * hidden_input->nb[2]); + for (int block_index = 0; block_index < config.block_count; ++block_index) { + const PIBlockWeights & block = weights.blocks[static_cast(block_index)]; + ggml_tensor * layer_hidden = + ggml_view_2d(graph_context, hidden_input, config.qwen_hidden_dim, static_cast(token_count), + hidden_input->nb[1], static_cast(block_index) * hidden_input->nb[2]); ggml_tensor * normalized = ada_norm(hidden, temb, block); - hidden = ggml_add( - graph_context, hidden, - attention(normalized, layer_hidden, block)); - ggml_tensor * ff = - ggml_norm(graph_context, hidden, - config.ada_norm_epsilon); - ff = linear(ff, block.feed_forward_input_weight, - block.feed_forward_input_bias); - ff = ggml_gelu(graph_context, ff); - ff = linear(ff, block.feed_forward_output_weight, - block.feed_forward_output_bias); - hidden = ggml_add(graph_context, hidden, ff); + hidden = ggml_add(graph_context, hidden, attention(normalized, layer_hidden, block)); + ggml_tensor * ff = ggml_norm(graph_context, hidden, config.ada_norm_epsilon); + ff = linear(ff, block.feed_forward_input_weight, block.feed_forward_input_bias); + ff = ggml_gelu(graph_context, ff); + ff = linear(ff, block.feed_forward_output_weight, block.feed_forward_output_bias); + hidden = ggml_add(graph_context, hidden, ff); } hidden = - ggml_relu(graph_context, - linear(hidden, weights.velocity_input_weight, - weights.velocity_input_bias)); - hidden = linear(hidden, weights.velocity_output_weight, - weights.velocity_output_bias); + ggml_relu(graph_context, linear(hidden, weights.velocity_input_weight, weights.velocity_input_bias)); + hidden = linear(hidden, weights.velocity_output_weight, weights.velocity_output_bias); ggml_tensor * velocity = ggml_view_2d( - graph_context, hidden, config.action_dim, config.horizon, - hidden->nb[1], - static_cast( - (include_state ? config.state_token_count : 0) + - config.future_token_count) * + graph_context, hidden, config.action_dim, config.horizon, hidden->nb[1], + static_cast((include_state ? config.state_token_count : 0) + config.future_token_count) * hidden->nb[1]); - actions = ggml_add( - graph_context, actions, - ggml_scale(graph_context, velocity, config.euler_dt)); + actions = ggml_add(graph_context, actions, ggml_scale(graph_context, velocity, config.euler_dt)); } output = actions; @@ -873,18 +530,15 @@ struct PIPolicy::Impl { ggml_set_output(output); graph = ggml_new_graph_custom(graph_context, kGraphSize, false); if (graph == nullptr) { - throw std::runtime_error( - "failed to create StarVLA PI graph"); + throw std::runtime_error("failed to create StarVLA PI graph"); } ggml_build_forward_expand(graph, output); ggml_backend_sched_reset(scheduler); if (!ggml_backend_sched_alloc_graph(scheduler, graph)) { - throw std::runtime_error( - "failed to allocate StarVLA PI graph"); + throw std::runtime_error("failed to allocate StarVLA PI graph"); } conditioning_token_count = token_count; - graph_uses_state = include_state; - ++graph_builds; + graph_uses_state = include_state; } }; @@ -892,8 +546,7 @@ PIPolicy::PIPolicy(std::unique_ptr impl) : impl_(std::move(impl)) {} PIPolicy::~PIPolicy() = default; -std::unique_ptr PIPolicy::load(const std::string & path, int n_threads, - int verbosity, std::string & error) { +std::unique_ptr PIPolicy::load(const std::string & path, int n_threads, int verbosity, std::string & error) { error.clear(); if (path.empty()) { error = "StarVLA PI policy path is required"; @@ -905,41 +558,36 @@ std::unique_ptr PIPolicy::load(const std::string & path, int n_threads impl->verbosity = verbosity; try { backend_scheduler_config scheduler_config; - scheduler_config.max_nodes = static_cast(kGraphSize); - scheduler_config.parallel = false; + scheduler_config.max_nodes = static_cast(kGraphSize); + scheduler_config.parallel = false; scheduler_config.op_offload = true; backend_loader backend; - if (!backend.load(impl->backend_cpu, impl->backends, impl->scheduler, - impl->buft_policy, true, scheduler_config, verbosity)) { + if (!backend.load(impl->backend_cpu, impl->backends, impl->scheduler, impl->buft_policy, true, scheduler_config, + verbosity)) { error = "failed to initialize StarVLA PI backend: " + backend.error(); return nullptr; } impl->mode = backend.mode(); PIGGUFLoader loader(impl->config, impl->weights); - if (!loader.load(path.c_str(), impl->buft_policy.model_buft, - impl->loaded, verbosity)) { + if (!loader.load(path.c_str(), impl->buft_policy.model_buft, impl->loaded, verbosity)) { error = loader.error(); return nullptr; } - if (impl->loaded.ctx_data == nullptr || - impl->loaded.model_buffer == nullptr) { + if (impl->loaded.ctx_data == nullptr || impl->loaded.model_buffer == nullptr) { error = "StarVLA PI policy GGUF has no tensors"; return nullptr; } - ggml_backend_buffer_set_usage( - impl->loaded.model_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + ggml_backend_buffer_set_usage(impl->loaded.model_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); impl->timestep_table = timestep_projection_table(impl->config); - impl->action_table = action_time_table(impl->config); + impl->action_table = action_time_table(impl->config); if (verbosity >= 1) { - std::fprintf( - stderr, - "%s: backend=%s qwen=%d width=%d blocks=%d horizon=%d " - "action_dim=%d profiles=%zu\n", - __func__, mode_name(impl->mode), impl->config.qwen_hidden_dim, - impl->config.dit_width, impl->config.block_count, - impl->config.horizon, impl->config.action_dim, - impl->config.normalization.profiles.size()); + std::fprintf(stderr, + "%s: backend=%s qwen=%d width=%d blocks=%d horizon=%d " + "action_dim=%d profiles=%zu\n", + __func__, backend_mode_name(impl->mode), impl->config.qwen_hidden_dim, impl->config.dit_width, + impl->config.block_count, impl->config.horizon, impl->config.action_dim, + impl->config.normalization.profiles.size()); } } catch (const std::exception & exception) { error = exception.what(); @@ -948,64 +596,48 @@ std::unique_ptr PIPolicy::load(const std::string & path, int n_threads return std::unique_ptr(new PIPolicy(std::move(impl))); } -bool PIPolicy::evaluate(const float * qwen_hidden_states, - size_t hidden_element_count, const float * state, - size_t state_element_count, const float * initial_noise, - size_t noise_element_count, - std::vector & normalized_actions, - std::string & error) { +bool PIPolicy::evaluate(const float * qwen_hidden_states, size_t hidden_element_count, const float * state, + size_t state_element_count, const float * initial_noise, size_t noise_element_count, + std::vector & normalized_actions, std::string & error) { normalized_actions.clear(); error.clear(); if (impl_ == nullptr || impl_->scheduler == nullptr) { error = "StarVLA PI policy is not initialized"; return false; } - const size_t layer_width = - static_cast(impl_->config.block_count) * - impl_->config.qwen_hidden_dim; - if (qwen_hidden_states == nullptr || layer_width == 0 || - hidden_element_count == 0 || + const size_t layer_width = static_cast(impl_->config.block_count) * impl_->config.qwen_hidden_dim; + if (qwen_hidden_states == nullptr || layer_width == 0 || hidden_element_count == 0 || hidden_element_count % layer_width != 0) { - error = - "StarVLA PI layer-wise Qwen conditioning tensor has an incompatible shape"; + error = "StarVLA PI layer-wise Qwen conditioning tensor has an incompatible shape"; return false; } const size_t token_count = hidden_element_count / layer_width; - if (token_count == 0 || - token_count > static_cast(std::numeric_limits::max())) { - error = - "StarVLA PI layer-wise Qwen conditioning tensor has an incompatible shape"; + if (token_count == 0 || token_count > static_cast(std::numeric_limits::max())) { + error = "StarVLA PI layer-wise Qwen conditioning tensor has an incompatible shape"; return false; } const bool include_state = state_element_count != 0; - if (include_state && - (state == nullptr || - state_element_count != static_cast(impl_->config.state_dim))) { + if (include_state && (state == nullptr || state_element_count != static_cast(impl_->config.state_dim))) { error = "StarVLA PI state tensor has an incompatible shape"; return false; } - const size_t expected_noise = - static_cast(impl_->config.horizon) * impl_->config.action_dim; + const size_t expected_noise = static_cast(impl_->config.horizon) * impl_->config.action_dim; if (initial_noise == nullptr || noise_element_count != expected_noise) { error = "StarVLA PI initial-noise tensor has an incompatible shape"; return false; } - if (std::any_of(qwen_hidden_states, - qwen_hidden_states + hidden_element_count, + if (std::any_of(qwen_hidden_states, qwen_hidden_states + hidden_element_count, [](float value) { return !std::isfinite(value); }) || (include_state && - std::any_of(state, state + state_element_count, - [](float value) { return !std::isfinite(value); })) || + std::any_of(state, state + state_element_count, [](float value) { return !std::isfinite(value); })) || std::any_of(initial_noise, initial_noise + noise_element_count, [](float value) { return !std::isfinite(value); })) { - error = - "StarVLA PI conditioning, state, and initial noise must be finite"; + error = "StarVLA PI conditioning, state, and initial noise must be finite"; return false; } try { - if (impl_->graph == nullptr || - impl_->conditioning_token_count != token_count || + if (impl_->graph == nullptr || impl_->conditioning_token_count != token_count || impl_->graph_uses_state != include_state) { impl_->build_graph(token_count, include_state); } @@ -1014,32 +646,25 @@ bool PIPolicy::evaluate(const float * qwen_hidden_states, return false; } - ggml_backend_tensor_set(impl_->hidden_input, qwen_hidden_states, 0, - hidden_element_count * sizeof(float)); + ggml_backend_tensor_set(impl_->hidden_input, qwen_hidden_states, 0, hidden_element_count * sizeof(float)); if (include_state) { - ggml_backend_tensor_set(impl_->state_input, state, 0, - state_element_count * sizeof(float)); + ggml_backend_tensor_set(impl_->state_input, state, 0, state_element_count * sizeof(float)); } - ggml_backend_tensor_set(impl_->noise_input, initial_noise, 0, - noise_element_count * sizeof(float)); - ggml_backend_tensor_set(impl_->timestep_projection_input, - impl_->timestep_table.data(), 0, + ggml_backend_tensor_set(impl_->noise_input, initial_noise, 0, noise_element_count * sizeof(float)); + ggml_backend_tensor_set(impl_->timestep_projection_input, impl_->timestep_table.data(), 0, impl_->timestep_table.size() * sizeof(float)); - ggml_backend_tensor_set(impl_->action_time_input, - impl_->action_table.data(), 0, + ggml_backend_tensor_set(impl_->action_time_input, impl_->action_table.data(), 0, impl_->action_table.size() * sizeof(float)); const float one = 1.0f; ggml_backend_tensor_set(impl_->scalar_one_input, &one, 0, sizeof(one)); set_backend_threads(impl_->backends, impl_->n_threads); - if (ggml_backend_sched_graph_compute(impl_->scheduler, impl_->graph) != - GGML_STATUS_SUCCESS) { + if (ggml_backend_sched_graph_compute(impl_->scheduler, impl_->graph) != GGML_STATUS_SUCCESS) { error = "StarVLA PI graph compute failed"; return false; } normalized_actions.resize(expected_noise); - ggml_backend_tensor_get(impl_->output, normalized_actions.data(), 0, - expected_noise * sizeof(float)); + ggml_backend_tensor_get(impl_->output, normalized_actions.data(), 0, expected_noise * sizeof(float)); if (std::any_of(normalized_actions.begin(), normalized_actions.end(), [](float value) { return !std::isfinite(value); })) { normalized_actions.clear(); @@ -1049,18 +674,15 @@ bool PIPolicy::evaluate(const float * qwen_hidden_states, return true; } -bool PIPolicy::unnormalize(const std::vector & normalized_actions, - const std::string & profile_key_value, - std::vector & actions, - std::string & error) const { +bool PIPolicy::unnormalize(const std::vector & normalized_actions, const std::string & profile_key_value, + std::vector & actions, std::string & error) const { if (impl_ == nullptr) { actions.clear(); error = "StarVLA PI policy is not initialized"; return false; } - return denormalize_actions(impl_->config.normalization, profile_key_value, - normalized_actions, impl_->config.horizon, - impl_->config.action_dim, actions, error); + return denormalize_actions(impl_->config.normalization, profile_key_value, normalized_actions, + impl_->config.horizon, impl_->config.action_dim, actions, error); } const PIPolicyConfig & PIPolicy::config() const { @@ -1071,11 +693,7 @@ const PIPolicyConfig & PIPolicy::config() const { } const char * PIPolicy::backend_name() const { - return impl_ != nullptr ? mode_name(impl_->mode) : "unknown"; -} - -size_t PIPolicy::graph_build_count() const { - return impl_ != nullptr ? impl_->graph_builds : 0; + return impl_ != nullptr ? backend_mode_name(impl_->mode) : "unknown"; } } // namespace robotcpp::starvla diff --git a/src/models/starvla/pi_policy.h b/src/models/starvla/pi_policy.h index 6f6923f..0b7b2a4 100644 --- a/src/models/starvla/pi_policy.h +++ b/src/models/starvla/pi_policy.h @@ -16,40 +16,40 @@ struct PIPolicyConfig { std::string text_filename; std::string mmproj_filename; - int qwen_hidden_dim = 0; + int qwen_hidden_dim = 0; int qwen_input_embedding_dim = 0; - int qwen_layer_count = 0; - int qwen_vocab_size = 0; + int qwen_layer_count = 0; + int qwen_vocab_size = 0; std::string cot_template; std::vector qwen_hidden_tuple_indices; int image_count = 0; std::vector image_names; - int image_framework_inference_pre_resize_width = 0; + int image_framework_inference_pre_resize_width = 0; int image_framework_inference_pre_resize_height = 0; - int image_processor_min_pixels = 0; - int image_processor_max_pixels = 0; - int image_patch_size = 0; - int image_spatial_merge_size = 0; - int image_min_token_count = 0; - int image_max_token_count = 0; + int image_processor_min_pixels = 0; + int image_processor_max_pixels = 0; + int image_patch_size = 0; + int image_spatial_merge_size = 0; + int image_min_token_count = 0; + int image_max_token_count = 0; - int dit_width = 0; - int block_count = 0; - int attention_head_count = 0; - int attention_head_dim = 0; - int cross_attention_dim = 0; - int feed_forward_dim = 0; - int mlp_hidden_dim = 0; - int state_dim = 0; - int action_dim = 0; - int horizon = 0; - int state_token_count = 0; - int future_token_count = 0; - int action_position_count = 0; + int dit_width = 0; + int block_count = 0; + int attention_head_count = 0; + int attention_head_dim = 0; + int cross_attention_dim = 0; + int feed_forward_dim = 0; + int mlp_hidden_dim = 0; + int state_dim = 0; + int action_dim = 0; + int horizon = 0; + int state_token_count = 0; + int future_token_count = 0; + int action_position_count = 0; int timestep_projection_dim = 0; int num_inference_timesteps = 0; - float ada_norm_epsilon = 0.0f; - float euler_dt = 0.0f; + float ada_norm_epsilon = 0.0f; + float euler_dt = 0.0f; std::vector timestep_ids; NormalizationConfig normalization; }; @@ -58,11 +58,10 @@ class PIPolicy { public: ~PIPolicy(); - PIPolicy(const PIPolicy &) = delete; + PIPolicy(const PIPolicy &) = delete; PIPolicy & operator=(const PIPolicy &) = delete; - static std::unique_ptr load(const std::string & path, int n_threads, - int verbosity, std::string & error); + static std::unique_ptr load(const std::string & path, int n_threads, int verbosity, std::string & error); // qwen_hidden_states is layer-major // [block_count, token_count, qwen_hidden_dim]. The legacy released @@ -70,17 +69,14 @@ class PIPolicy { // head. state is either omitted (the official Bridge deployment path) or // one token [state_dim], and initial_noise is token-major // [horizon, action_dim]. - bool evaluate(const float * qwen_hidden_states, size_t hidden_element_count, - const float * state, size_t state_element_count, - const float * initial_noise, size_t noise_element_count, + bool evaluate(const float * qwen_hidden_states, size_t hidden_element_count, const float * state, + size_t state_element_count, const float * initial_noise, size_t noise_element_count, std::vector & normalized_actions, std::string & error); - bool unnormalize(const std::vector & normalized_actions, - const std::string & profile_key, std::vector & actions, - std::string & error) const; + bool unnormalize(const std::vector & normalized_actions, const std::string & profile_key, + std::vector & actions, std::string & error) const; const PIPolicyConfig & config() const; const char * backend_name() const; - size_t graph_build_count() const; private: struct Impl; diff --git a/src/models/starvla/pi_v3_policy.cpp b/src/models/starvla/pi_v3_policy.cpp index 8cc1fd1..bedc679 100644 --- a/src/models/starvla/pi_v3_policy.cpp +++ b/src/models/starvla/pi_v3_policy.cpp @@ -5,10 +5,10 @@ #include "gguf.h" #include "models/ggml_backend.h" #include "models/gguf_loader.h" +#include "models/starvla/policy_gguf.h" #include #include #include -#include #include #include #include @@ -19,163 +19,60 @@ namespace robotcpp::starvla { namespace { -constexpr size_t kGraphSize = 32768; -constexpr int kKQMaskPad = 32; +constexpr size_t kGraphSize = 32768; +constexpr int kKQMaskPad = 32; constexpr int kReleasedLayerCount = 36; struct PIV3BlockWeights { - ggml_tensor * ada_norm_weight = nullptr; - ggml_tensor * ada_norm_bias = nullptr; - ggml_tensor * query_weight = nullptr; - ggml_tensor * query_bias = nullptr; - ggml_tensor * key_weight = nullptr; - ggml_tensor * key_bias = nullptr; - ggml_tensor * value_weight = nullptr; - ggml_tensor * value_bias = nullptr; - ggml_tensor * attention_output_weight = nullptr; - ggml_tensor * attention_output_bias = nullptr; - ggml_tensor * feed_forward_input_weight = nullptr; - ggml_tensor * feed_forward_input_bias = nullptr; + ggml_tensor * ada_norm_weight = nullptr; + ggml_tensor * ada_norm_bias = nullptr; + ggml_tensor * query_weight = nullptr; + ggml_tensor * query_bias = nullptr; + ggml_tensor * key_weight = nullptr; + ggml_tensor * key_bias = nullptr; + ggml_tensor * value_weight = nullptr; + ggml_tensor * value_bias = nullptr; + ggml_tensor * attention_output_weight = nullptr; + ggml_tensor * attention_output_bias = nullptr; + ggml_tensor * feed_forward_input_weight = nullptr; + ggml_tensor * feed_forward_input_bias = nullptr; ggml_tensor * feed_forward_output_weight = nullptr; - ggml_tensor * feed_forward_output_bias = nullptr; + ggml_tensor * feed_forward_output_bias = nullptr; }; struct PIV3ProjectorWeights { - ggml_tensor * norm_weight = nullptr; - ggml_tensor * norm_bias = nullptr; + ggml_tensor * norm_weight = nullptr; + ggml_tensor * norm_bias = nullptr; ggml_tensor * projection_weight = nullptr; - ggml_tensor * projection_bias = nullptr; + ggml_tensor * projection_bias = nullptr; }; struct PIV3Weights { - ggml_tensor * timestep_input_weight = nullptr; - ggml_tensor * timestep_input_bias = nullptr; + ggml_tensor * timestep_input_weight = nullptr; + ggml_tensor * timestep_input_bias = nullptr; ggml_tensor * timestep_output_weight = nullptr; - ggml_tensor * timestep_output_bias = nullptr; + ggml_tensor * timestep_output_bias = nullptr; std::vector blocks; std::vector projectors; - ggml_tensor * action_input_weight = nullptr; - ggml_tensor * action_input_bias = nullptr; + ggml_tensor * action_input_weight = nullptr; + ggml_tensor * action_input_bias = nullptr; ggml_tensor * action_time_mix_weight = nullptr; - ggml_tensor * action_time_mix_bias = nullptr; - ggml_tensor * action_output_weight = nullptr; - ggml_tensor * action_output_bias = nullptr; - ggml_tensor * velocity_input_weight = nullptr; - ggml_tensor * velocity_input_bias = nullptr; + ggml_tensor * action_time_mix_bias = nullptr; + ggml_tensor * action_output_weight = nullptr; + ggml_tensor * action_output_bias = nullptr; + ggml_tensor * velocity_input_weight = nullptr; + ggml_tensor * velocity_input_bias = nullptr; ggml_tensor * velocity_output_weight = nullptr; - ggml_tensor * velocity_output_bias = nullptr; - ggml_tensor * future_tokens = nullptr; - ggml_tensor * action_position = nullptr; + ggml_tensor * velocity_output_bias = nullptr; + ggml_tensor * future_tokens = nullptr; + ggml_tensor * action_position = nullptr; }; -int require_key(gguf_context * gguf, const char * key, gguf_type type) { - const int index = gguf_find_key(gguf, key); - if (index < 0) { - throw std::runtime_error(std::string("missing required StarVLA GGUF metadata: ") + key); - } - if (gguf_get_kv_type(gguf, index) != type) { - throw std::runtime_error(std::string("invalid StarVLA GGUF metadata type: ") + key); - } - return index; -} - -std::string require_string(gguf_context * gguf, const char * key) { - return gguf_get_val_str(gguf, require_key(gguf, key, GGUF_TYPE_STRING)); -} - -int require_i32(gguf_context * gguf, const char * key) { - return gguf_get_val_i32(gguf, require_key(gguf, key, GGUF_TYPE_INT32)); -} - -float require_f32(gguf_context * gguf, const char * key) { - return gguf_get_val_f32(gguf, require_key(gguf, key, GGUF_TYPE_FLOAT32)); -} - -bool require_bool(gguf_context * gguf, const char * key) { - return gguf_get_val_bool(gguf, require_key(gguf, key, GGUF_TYPE_BOOL)); -} - -int require_array(gguf_context * gguf, const char * key, gguf_type element_type) { - const int index = require_key(gguf, key, GGUF_TYPE_ARRAY); - if (gguf_get_arr_type(gguf, index) != element_type) { - throw std::runtime_error(std::string("invalid StarVLA GGUF array element type: ") + key); - } - return index; -} - -std::vector require_string_array(gguf_context * gguf, const char * key) { - const int index = require_array(gguf, key, GGUF_TYPE_STRING); - const size_t count = gguf_get_arr_n(gguf, index); - std::vector result; - result.reserve(count); - for (size_t i = 0; i < count; ++i) { - result.emplace_back(gguf_get_arr_str(gguf, index, i)); - } - return result; -} - -std::vector require_i32_array(gguf_context * gguf, const char * key) { - const int index = require_array(gguf, key, GGUF_TYPE_INT32); - const size_t count = gguf_get_arr_n(gguf, index); - const auto * data = static_cast(gguf_get_arr_data(gguf, index)); - if (data == nullptr && count != 0) { - throw std::runtime_error(std::string("missing StarVLA GGUF array data: ") + key); - } - return count == 0 ? std::vector() : std::vector(data, data + count); -} - -std::vector require_f32_array(gguf_context * gguf, const char * key) { - const int index = require_array(gguf, key, GGUF_TYPE_FLOAT32); - const size_t count = gguf_get_arr_n(gguf, index); - const auto * data = static_cast(gguf_get_arr_data(gguf, index)); - if (data == nullptr && count != 0) { - throw std::runtime_error(std::string("missing StarVLA GGUF array data: ") + key); - } - return count == 0 ? std::vector() : std::vector(data, data + count); -} - -std::vector require_bool_array(gguf_context * gguf, const char * key) { - const int index = require_array(gguf, key, GGUF_TYPE_BOOL); - const size_t count = gguf_get_arr_n(gguf, index); - const auto * data = static_cast(gguf_get_arr_data(gguf, index)); - if (data == nullptr && count != 0) { - throw std::runtime_error(std::string("missing StarVLA GGUF array data: ") + key); - } - std::vector result(count); - for (size_t i = 0; i < count; ++i) { - result[i] = data[i] != 0 ? 1 : 0; - } - return result; -} - -std::string profile_key(int profile_index, const char * suffix) { - return "starvla.normalization.profile." + std::to_string(profile_index) + "." + suffix; -} - -bool has_shape(const ggml_tensor * tensor, std::initializer_list expected) { - if (tensor == nullptr || static_cast(ggml_n_dims(tensor)) != expected.size()) { - return false; - } - size_t dimension = 0; - for (const int64_t value : expected) { - if (tensor->ne[dimension++] != value) { - return false; - } - } - return true; -} - -const char * mode_name(backend_mode mode) { - switch (mode) { - case backend_mode::cpu: - return "cpu"; - case backend_mode::cuda: - return "cuda"; - case backend_mode::metal: - return "metal"; - } - return "unknown"; -} +using detail::has_shape; +using detail::require_f32; +using detail::require_i32; +using detail::require_string; +using detail::require_string_array; std::vector integer_range(int first, int count) { std::vector result(static_cast(count)); @@ -192,8 +89,7 @@ class PIV3GGUFLoader final : public gguf_loader { protected: bool parse_metadata(gguf_context * gguf) override { if (require_string(gguf, "general.architecture") != "starvla-policy" || - require_i32(gguf, "starvla.schema_version") != 1 || - require_string(gguf, "starvla.framework") != "pi_v3") { + require_i32(gguf, "starvla.schema_version") != 1 || require_string(gguf, "starvla.framework") != "pi_v3") { throw std::runtime_error("GGUF is not a supported StarVLA PI-v3 policy"); } @@ -201,66 +97,46 @@ class PIV3GGUFLoader final : public gguf_loader { if (config_.backbone_arch != "qwen3_vl") { throw std::runtime_error("StarVLA PI-v3 requires a Qwen3-VL backbone"); } - config_.bundle_uuid = require_string(gguf, "starvla.bundle.uuid"); - config_.text_filename = require_string(gguf, "starvla.component.text.filename"); + config_.bundle_uuid = require_string(gguf, "starvla.bundle.uuid"); + config_.text_filename = require_string(gguf, "starvla.component.text.filename"); config_.mmproj_filename = require_string(gguf, "starvla.component.mmproj.filename"); - if (config_.bundle_uuid.empty() || config_.text_filename.empty() || - config_.mmproj_filename.empty()) { + if (config_.bundle_uuid.empty() || config_.text_filename.empty() || config_.mmproj_filename.empty()) { throw std::runtime_error("StarVLA PI-v3 bundle metadata is incomplete"); } - config_.qwen_hidden_dim = require_i32(gguf, "starvla.qwen.hidden_size"); - const int embedding_key = gguf_find_key(gguf, "starvla.qwen.input_embedding_size"); - config_.qwen_input_embedding_dim = - embedding_key < 0 ? 4 * config_.qwen_hidden_dim - : require_i32(gguf, "starvla.qwen.input_embedding_size"); - config_.qwen_layer_count = require_i32(gguf, "starvla.qwen.layer_count"); - config_.qwen_vocab_size = require_i32(gguf, "starvla.qwen.vocab_size"); - config_.cot_template = require_string(gguf, "starvla.prompt.cot_template"); - - config_.image_count = require_i32(gguf, "starvla.image.count"); - config_.image_names = require_string_array(gguf, "starvla.image.names"); - config_.image_processor_min_pixels = - require_i32(gguf, "starvla.image.processor_min_pixels"); - config_.image_processor_max_pixels = - require_i32(gguf, "starvla.image.processor_max_pixels"); - config_.image_patch_size = require_i32(gguf, "starvla.image.patch_size"); - config_.image_spatial_merge_size = - require_i32(gguf, "starvla.image.spatial_merge_size"); - config_.image_min_token_count = - require_i32(gguf, "starvla.image.min_token_count"); - config_.image_max_token_count = - require_i32(gguf, "starvla.image.max_token_count"); - - config_.dit_width = require_i32(gguf, "starvla.pi_v3.dit_width"); - config_.block_count = require_i32(gguf, "starvla.pi_v3.block_count"); - config_.projector_count = require_i32(gguf, "starvla.pi_v3.projector_count"); - config_.attention_head_count = - require_i32(gguf, "starvla.pi_v3.attention_head_count"); - config_.attention_head_dim = - require_i32(gguf, "starvla.pi_v3.attention_head_dim"); - config_.feed_forward_dim = - require_i32(gguf, "starvla.pi_v3.feed_forward_dim"); - config_.mlp_hidden_dim = - require_i32(gguf, "starvla.pi_v3.mlp_hidden_dimension"); - config_.future_token_count = - require_i32(gguf, "starvla.pi_v3.future_token_count"); - config_.action_position_count = - require_i32(gguf, "starvla.pi_v3.action_position_count"); - config_.no_state_sequence_length = - require_i32(gguf, "starvla.pi_v3.no_state_sequence_length"); - config_.timestep_projection_dim = - require_i32(gguf, "starvla.pi_v3.timestep_projection_dim"); - config_.num_timestep_buckets = - require_i32(gguf, "starvla.pi_v3.num_timestep_buckets"); - config_.num_inference_timesteps = - require_i32(gguf, "starvla.pi_v3.num_inference_timesteps"); - config_.ada_norm_epsilon = require_f32(gguf, "starvla.pi_v3.ada_norm_epsilon"); - config_.projector_norm_epsilon = - require_f32(gguf, "starvla.pi_v3.projector_norm_epsilon"); - config_.euler_dt = require_f32(gguf, "starvla.pi_v3.euler_dt"); - config_.action_dim = require_i32(gguf, "starvla.action.dimension"); - config_.horizon = require_i32(gguf, "starvla.action.horizon"); + config_.qwen_hidden_dim = require_i32(gguf, "starvla.qwen.hidden_size"); + config_.qwen_input_embedding_dim = require_i32(gguf, "starvla.qwen.input_embedding_size"); + config_.qwen_layer_count = require_i32(gguf, "starvla.qwen.layer_count"); + config_.qwen_vocab_size = require_i32(gguf, "starvla.qwen.vocab_size"); + config_.cot_template = require_string(gguf, "starvla.prompt.cot_template"); + + config_.image_count = require_i32(gguf, "starvla.image.count"); + config_.image_names = require_string_array(gguf, "starvla.image.names"); + config_.image_processor_min_pixels = require_i32(gguf, "starvla.image.processor_min_pixels"); + config_.image_processor_max_pixels = require_i32(gguf, "starvla.image.processor_max_pixels"); + config_.image_patch_size = require_i32(gguf, "starvla.image.patch_size"); + config_.image_spatial_merge_size = require_i32(gguf, "starvla.image.spatial_merge_size"); + config_.image_min_token_count = require_i32(gguf, "starvla.image.min_token_count"); + config_.image_max_token_count = require_i32(gguf, "starvla.image.max_token_count"); + + config_.dit_width = require_i32(gguf, "starvla.pi_v3.dit_width"); + config_.block_count = require_i32(gguf, "starvla.pi_v3.block_count"); + config_.projector_count = require_i32(gguf, "starvla.pi_v3.projector_count"); + config_.attention_head_count = require_i32(gguf, "starvla.pi_v3.attention_head_count"); + config_.attention_head_dim = require_i32(gguf, "starvla.pi_v3.attention_head_dim"); + config_.feed_forward_dim = require_i32(gguf, "starvla.pi_v3.feed_forward_dim"); + config_.mlp_hidden_dim = require_i32(gguf, "starvla.pi_v3.mlp_hidden_dimension"); + config_.future_token_count = require_i32(gguf, "starvla.pi_v3.future_token_count"); + config_.action_position_count = require_i32(gguf, "starvla.pi_v3.action_position_count"); + config_.no_state_sequence_length = require_i32(gguf, "starvla.pi_v3.no_state_sequence_length"); + config_.timestep_projection_dim = require_i32(gguf, "starvla.pi_v3.timestep_projection_dim"); + config_.num_timestep_buckets = require_i32(gguf, "starvla.pi_v3.num_timestep_buckets"); + config_.num_inference_timesteps = require_i32(gguf, "starvla.pi_v3.num_inference_timesteps"); + config_.ada_norm_epsilon = require_f32(gguf, "starvla.pi_v3.ada_norm_epsilon"); + config_.projector_norm_epsilon = require_f32(gguf, "starvla.pi_v3.projector_norm_epsilon"); + config_.euler_dt = require_f32(gguf, "starvla.pi_v3.euler_dt"); + config_.action_dim = require_i32(gguf, "starvla.action.dimension"); + config_.horizon = require_i32(gguf, "starvla.action.horizon"); const bool valid = config_.qwen_hidden_dim > 0 && config_.qwen_input_embedding_dim > 0 && @@ -268,25 +144,18 @@ class PIV3GGUFLoader final : public gguf_loader { !config_.cot_template.empty() && config_.image_count > 0 && config_.image_names.size() == static_cast(config_.image_count) && config_.image_processor_min_pixels > 0 && - config_.image_processor_max_pixels >= config_.image_processor_min_pixels && - config_.image_patch_size > 0 && config_.image_spatial_merge_size > 0 && - config_.image_min_token_count > 0 && - config_.image_max_token_count >= config_.image_min_token_count && - config_.dit_width > 0 && config_.block_count == kReleasedLayerCount && - config_.projector_count == config_.block_count && + config_.image_processor_max_pixels >= config_.image_processor_min_pixels && config_.image_patch_size > 0 && + config_.image_spatial_merge_size > 0 && config_.image_min_token_count > 0 && + config_.image_max_token_count >= config_.image_min_token_count && config_.dit_width > 0 && + config_.block_count == kReleasedLayerCount && config_.projector_count == config_.block_count && config_.attention_head_count > 0 && config_.attention_head_dim > 0 && config_.attention_head_count * config_.attention_head_dim == config_.dit_width && - config_.feed_forward_dim > 0 && config_.mlp_hidden_dim > 0 && - config_.action_dim > 0 && config_.horizon > 0 && - config_.future_token_count > 0 && - config_.action_position_count >= config_.horizon && - config_.no_state_sequence_length == - config_.future_token_count + config_.horizon && - config_.timestep_projection_dim >= 4 && - config_.timestep_projection_dim % 2 == 0 && + config_.feed_forward_dim > 0 && config_.mlp_hidden_dim > 0 && config_.action_dim > 0 && + config_.horizon > 0 && config_.future_token_count > 0 && config_.action_position_count >= config_.horizon && + config_.no_state_sequence_length == config_.future_token_count + config_.horizon && + config_.timestep_projection_dim >= 4 && config_.timestep_projection_dim % 2 == 0 && config_.num_timestep_buckets > 0 && config_.num_inference_timesteps == 4 && - config_.ada_norm_epsilon > 0.0f && config_.projector_norm_epsilon > 0.0f && - config_.euler_dt > 0.0f; + config_.ada_norm_epsilon > 0.0f && config_.projector_norm_epsilon > 0.0f && config_.euler_dt > 0.0f; if (!valid) { throw std::runtime_error("StarVLA PI-v3 metadata has incompatible dimensions"); } @@ -298,44 +167,7 @@ class PIV3GGUFLoader final : public gguf_loader { step * config_.num_timestep_buckets / config_.num_inference_timesteps; } - NormalizationConfig & normalization = config_.normalization; - normalization.clip_actions = require_bool(gguf, "starvla.normalization.clip_actions"); - normalization.binary_threshold = - require_f32(gguf, "starvla.normalization.binary_threshold"); - normalization.binary_comparison = - require_string(gguf, "starvla.normalization.binary_comparison"); - normalization.continuous_dimensions = - require_i32_array(gguf, "starvla.action.continuous_dimensions"); - normalization.binary_dimensions = - require_i32_array(gguf, "starvla.action.binary_dimensions"); - const int profile_count = - require_i32(gguf, "starvla.normalization.profile_count"); - const std::vector keys = - require_string_array(gguf, "starvla.normalization.profile_keys"); - if (profile_count <= 0 || keys.size() != static_cast(profile_count)) { - throw std::runtime_error("StarVLA PI-v3 normalization profiles are inconsistent"); - } - normalization.profiles.clear(); - normalization.profiles.reserve(static_cast(profile_count)); - for (int index = 0; index < profile_count; ++index) { - NormalizationProfile profile; - profile.key = require_string(gguf, profile_key(index, "key").c_str()); - profile.action_q01 = - require_f32_array(gguf, profile_key(index, "action_q01").c_str()); - profile.action_q99 = - require_f32_array(gguf, profile_key(index, "action_q99").c_str()); - profile.action_mask = - require_bool_array(gguf, profile_key(index, "action_mask").c_str()); - if (profile.key != keys[static_cast(index)]) { - throw std::runtime_error("StarVLA PI-v3 normalization profile order is inconsistent"); - } - normalization.profiles.push_back(std::move(profile)); - } - std::string normalization_error; - if (!validate_normalization_config(normalization, config_.action_dim, - normalization_error)) { - throw std::runtime_error(normalization_error); - } + config_.normalization = detail::require_normalization(gguf, config_.action_dim); return true; } @@ -371,8 +203,7 @@ class PIV3GGUFLoader final : public gguf_loader { weights_.projectors.clear(); weights_.projectors.reserve(static_cast(config_.projector_count)); for (int projector = 0; projector < config_.projector_count; ++projector) { - const std::string prefix = "starvla.policy.pi_v3.projector." + - std::to_string(projector) + "."; + const std::string prefix = "starvla.policy.pi_v3.projector." + std::to_string(projector) + "."; PIV3ProjectorWeights current; bind(current.norm_weight, prefix + "norm.weight"); bind(current.norm_bias, prefix + "norm.bias"); @@ -414,12 +245,9 @@ class PIV3GGUFLoader final : public gguf_loader { } for (const PIV3BlockWeights & current : weights_.blocks) { if (!has_shape(current.ada_norm_weight, {width, 2 * width}) || - !has_shape(current.ada_norm_bias, {2 * width}) || - !has_shape(current.query_weight, {width, width}) || - !has_shape(current.query_bias, {width}) || - !has_shape(current.key_weight, {width, width}) || - !has_shape(current.key_bias, {width}) || - !has_shape(current.value_weight, {width, width}) || + !has_shape(current.ada_norm_bias, {2 * width}) || !has_shape(current.query_weight, {width, width}) || + !has_shape(current.query_bias, {width}) || !has_shape(current.key_weight, {width, width}) || + !has_shape(current.key_bias, {width}) || !has_shape(current.value_weight, {width, width}) || !has_shape(current.value_bias, {width}) || !has_shape(current.attention_output_weight, {width, width}) || !has_shape(current.attention_output_bias, {width}) || @@ -447,36 +275,36 @@ class PIV3GGUFLoader final : public gguf_loader { }; std::vector timestep_projection_table(const PIV3PolicyConfig & config) { - const int dim = config.timestep_projection_dim; - const int half = dim / 2; + const int dim = config.timestep_projection_dim; + const int half = dim / 2; const float denominator = static_cast(half - 1); std::vector table(static_cast(dim) * 4, 0.0f); for (int step = 0; step < 4; ++step) { const float timestep = static_cast(config.timestep_ids[static_cast(step)]); - float * row = table.data() + static_cast(step) * dim; + float * row = table.data() + static_cast(step) * dim; for (int index = 0; index < half; ++index) { const float frequency = std::exp(-std::log(10000.0f) * static_cast(index) / denominator); - const float angle = timestep * frequency; - row[index] = std::cos(angle); - row[index + half] = std::sin(angle); + const float angle = timestep * frequency; + row[index] = std::cos(angle); + row[index + half] = std::sin(angle); } } return table; } std::vector action_time_table(const PIV3PolicyConfig & config) { - const int dim = config.dit_width; - const int half = dim / 2; + const int dim = config.dit_width; + const int half = dim / 2; const float denominator = static_cast(half); std::vector table(static_cast(dim) * 4, 0.0f); for (int step = 0; step < 4; ++step) { const float timestep = static_cast(config.timestep_ids[static_cast(step)]); - float * row = table.data() + static_cast(step) * dim; + float * row = table.data() + static_cast(step) * dim; for (int index = 0; index < half; ++index) { const float frequency = std::exp(-std::log(10000.0f) * static_cast(index) / denominator); - const float angle = timestep * frequency; - row[index] = std::sin(angle); - row[index + half] = std::cos(angle); + const float angle = timestep * frequency; + row[index] = std::sin(angle); + row[index + half] = std::cos(angle); } } return table; @@ -492,19 +320,19 @@ struct PIV3Policy::Impl { std::vector backends; ggml_backend_sched_t scheduler = nullptr; backend_buft_policy buft_policy; - backend_mode mode = backend_mode::cpu; - int n_threads = 0; - int verbosity = 0; - ggml_context * graph_context = nullptr; - ggml_cgraph * graph = nullptr; - ggml_tensor * hidden_input = nullptr; - ggml_tensor * cross_mask_input = nullptr; - ggml_tensor * noise_input = nullptr; + backend_mode mode = backend_mode::cpu; + int n_threads = 0; + int verbosity = 0; + ggml_context * graph_context = nullptr; + ggml_cgraph * graph = nullptr; + ggml_tensor * hidden_input = nullptr; + ggml_tensor * cross_mask_input = nullptr; + ggml_tensor * noise_input = nullptr; ggml_tensor * timestep_projection_input = nullptr; - ggml_tensor * action_time_input = nullptr; - ggml_tensor * scalar_one_input = nullptr; - ggml_tensor * output = nullptr; - size_t conditioning_token_count = 0; + ggml_tensor * action_time_input = nullptr; + ggml_tensor * scalar_one_input = nullptr; + ggml_tensor * output = nullptr; + size_t conditioning_token_count = 0; std::vector timestep_table; std::vector action_table; ~Impl() { @@ -544,15 +372,15 @@ struct PIV3Policy::Impl { ggml_free(graph_context); graph_context = nullptr; } - graph = nullptr; - hidden_input = nullptr; - cross_mask_input = nullptr; - noise_input = nullptr; + graph = nullptr; + hidden_input = nullptr; + cross_mask_input = nullptr; + noise_input = nullptr; timestep_projection_input = nullptr; - action_time_input = nullptr; - scalar_one_input = nullptr; - output = nullptr; - conditioning_token_count = 0; + action_time_input = nullptr; + scalar_one_input = nullptr; + output = nullptr; + conditioning_token_count = 0; } void build_graph(size_t token_count) { @@ -562,34 +390,30 @@ struct PIV3Policy::Impl { } ggml_init_params params{}; - params.mem_size = kGraphSize * ggml_tensor_overhead() + - ggml_graph_overhead_custom(kGraphSize, false); + params.mem_size = kGraphSize * ggml_tensor_overhead() + ggml_graph_overhead_custom(kGraphSize, false); params.mem_buffer = nullptr; - params.no_alloc = true; - graph_context = ggml_init(params); + params.no_alloc = true; + graph_context = ggml_init(params); if (graph_context == nullptr) { throw std::runtime_error("failed to initialize StarVLA PI_v3 graph context"); } - const int width = config.dit_width; - const int heads = config.attention_head_count; - const int head_dim = config.attention_head_dim; + const int width = config.dit_width; + const int heads = config.attention_head_count; + const int head_dim = config.attention_head_dim; const int sequence_length = config.no_state_sequence_length; - const int mask_queries = GGML_PAD(sequence_length, kKQMaskPad); + const int mask_queries = GGML_PAD(sequence_length, kKQMaskPad); hidden_input = ggml_new_tensor_3d(graph_context, GGML_TYPE_F32, config.qwen_hidden_dim, static_cast(token_count), config.qwen_layer_count); - cross_mask_input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, - static_cast(token_count), mask_queries); - noise_input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, - config.action_dim, config.horizon); - timestep_projection_input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, - config.timestep_projection_dim, 4); - action_time_input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, width, 4); - scalar_one_input = ggml_new_tensor_1d(graph_context, GGML_TYPE_F32, 1); + cross_mask_input = + ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, static_cast(token_count), mask_queries); + noise_input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, config.action_dim, config.horizon); + timestep_projection_input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, config.timestep_projection_dim, 4); + action_time_input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, width, 4); + scalar_one_input = ggml_new_tensor_1d(graph_context, GGML_TYPE_F32, 1); if (hidden_input == nullptr || cross_mask_input == nullptr || noise_input == nullptr || - timestep_projection_input == nullptr || action_time_input == nullptr || - scalar_one_input == nullptr) { + timestep_projection_input == nullptr || action_time_input == nullptr || scalar_one_input == nullptr) { throw std::runtime_error("failed to create StarVLA PI_v3 graph inputs"); } ggml_set_name(hidden_input, "starvla_pi_v3_qwen_hidden_states"); @@ -606,59 +430,52 @@ struct PIV3Policy::Impl { ggml_set_input(scalar_one_input); auto f32 = [&](ggml_tensor * tensor) { - return tensor->type == GGML_TYPE_F32 ? tensor : - ggml_cast(graph_context, tensor, GGML_TYPE_F32); + return tensor->type == GGML_TYPE_F32 ? tensor : ggml_cast(graph_context, tensor, GGML_TYPE_F32); }; auto bf16_roundtrip = [&](ggml_tensor * tensor) { - return ggml_cast(graph_context, - ggml_cast(graph_context, tensor, GGML_TYPE_BF16), - GGML_TYPE_F32); + return ggml_cast(graph_context, ggml_cast(graph_context, tensor, GGML_TYPE_BF16), GGML_TYPE_F32); }; auto linear = [&](ggml_tensor * value, ggml_tensor * weight, ggml_tensor * bias) { ggml_tensor * projected = ggml_mul_mat(graph_context, weight, value); ggml_mul_mat_set_prec(projected, GGML_PREC_F32); return ggml_add(graph_context, projected, f32(bias)); }; - auto projector_linear = [&](ggml_tensor * value, ggml_tensor * weight, - ggml_tensor * bias) { - ggml_tensor * bf16_value = ggml_cast(graph_context, value, GGML_TYPE_BF16); + auto projector_linear = [&](ggml_tensor * value, ggml_tensor * weight, ggml_tensor * bias) { + ggml_tensor * bf16_value = ggml_cast(graph_context, value, GGML_TYPE_BF16); ggml_tensor * bf16_weight = ggml_cast(graph_context, weight, GGML_TYPE_BF16); - ggml_tensor * projected = - ggml_mul_mat(graph_context, bf16_weight, bf16_value); + ggml_tensor * projected = ggml_mul_mat(graph_context, bf16_weight, bf16_value); ggml_mul_mat_set_prec(projected, GGML_PREC_F32); projected = ggml_add(graph_context, projected, bf16_roundtrip(bias)); return bf16_roundtrip(projected); }; - auto ada_norm = [&](ggml_tensor * value, ggml_tensor * temb, - const PIV3BlockWeights & block) { - ggml_tensor * modulation = linear(ggml_silu(graph_context, temb), - block.ada_norm_weight, block.ada_norm_bias); + auto ada_norm = [&](ggml_tensor * value, ggml_tensor * temb, const PIV3BlockWeights & block) { + ggml_tensor * modulation = + linear(ggml_silu(graph_context, temb), block.ada_norm_weight, block.ada_norm_bias); ggml_tensor * scale = ggml_view_1d(graph_context, modulation, width, 0); - ggml_tensor * shift = ggml_view_1d(graph_context, modulation, width, - static_cast(width) * sizeof(float)); + ggml_tensor * shift = + ggml_view_1d(graph_context, modulation, width, static_cast(width) * sizeof(float)); ggml_tensor * normalized = ggml_norm(graph_context, value, config.ada_norm_epsilon); return ggml_add(graph_context, - ggml_mul(graph_context, normalized, - ggml_add(graph_context, scale, scalar_one_input)), + ggml_mul(graph_context, normalized, ggml_add(graph_context, scale, scalar_one_input)), shift); }; auto attention = [&](ggml_tensor * query_source, ggml_tensor * key_value_source, const PIV3BlockWeights & block) { - const int64_t query_count = query_source->ne[1]; + const int64_t query_count = query_source->ne[1]; const int64_t key_value_count = key_value_source->ne[1]; - ggml_tensor * query = linear(query_source, block.query_weight, block.query_bias); - ggml_tensor * key = linear(key_value_source, block.key_weight, block.key_bias); - ggml_tensor * value = linear(key_value_source, block.value_weight, block.value_bias); - query = ggml_reshape_3d(graph_context, query, head_dim, heads, query_count); - key = ggml_reshape_3d(graph_context, key, head_dim, heads, key_value_count); - value = ggml_reshape_3d(graph_context, value, head_dim, heads, key_value_count); - query = ggml_permute(graph_context, query, 0, 2, 1, 3); - key = ggml_permute(graph_context, key, 0, 2, 1, 3); - value = ggml_cont(graph_context, ggml_permute(graph_context, value, 1, 2, 0, 3)); - ggml_tensor * scores = ggml_mul_mat(graph_context, key, query); + ggml_tensor * query = linear(query_source, block.query_weight, block.query_bias); + ggml_tensor * key = linear(key_value_source, block.key_weight, block.key_bias); + ggml_tensor * value = linear(key_value_source, block.value_weight, block.value_bias); + query = ggml_reshape_3d(graph_context, query, head_dim, heads, query_count); + key = ggml_reshape_3d(graph_context, key, head_dim, heads, key_value_count); + value = ggml_reshape_3d(graph_context, value, head_dim, heads, key_value_count); + query = ggml_permute(graph_context, query, 0, 2, 1, 3); + key = ggml_permute(graph_context, key, 0, 2, 1, 3); + value = ggml_cont(graph_context, ggml_permute(graph_context, value, 1, 2, 0, 3)); + ggml_tensor * scores = ggml_mul_mat(graph_context, key, query); ggml_mul_mat_set_prec(scores, GGML_PREC_F32); - scores = ggml_soft_max_ext(graph_context, scores, cross_mask_input, - 1.0f / std::sqrt(static_cast(head_dim)), 0.0f); + scores = ggml_soft_max_ext(graph_context, scores, cross_mask_input, + 1.0f / std::sqrt(static_cast(head_dim)), 0.0f); ggml_tensor * attended = ggml_mul_mat(graph_context, value, scores); ggml_mul_mat_set_prec(attended, GGML_PREC_F32); attended = ggml_permute(graph_context, attended, 0, 2, 1, 3); @@ -670,78 +487,68 @@ struct PIV3Policy::Impl { projected_hidden_states.reserve(static_cast(config.projector_count)); for (int layer = 0; layer < config.projector_count; ++layer) { const PIV3ProjectorWeights & projector = weights.projectors[static_cast(layer)]; - ggml_tensor * layer_hidden = ggml_view_2d( - graph_context, hidden_input, config.qwen_hidden_dim, - static_cast(token_count), hidden_input->nb[1], - static_cast(layer) * hidden_input->nb[2]); + ggml_tensor * layer_hidden = + ggml_view_2d(graph_context, hidden_input, config.qwen_hidden_dim, static_cast(token_count), + hidden_input->nb[1], static_cast(layer) * hidden_input->nb[2]); layer_hidden = bf16_roundtrip(layer_hidden); layer_hidden = ggml_norm(graph_context, layer_hidden, config.projector_norm_epsilon); layer_hidden = ggml_mul(graph_context, layer_hidden, f32(projector.norm_weight)); layer_hidden = ggml_add(graph_context, layer_hidden, f32(projector.norm_bias)); - ggml_tensor * projected = projector_linear( - layer_hidden, projector.projection_weight, - projector.projection_bias); + ggml_tensor * projected = + projector_linear(layer_hidden, projector.projection_weight, projector.projection_bias); projected_hidden_states.push_back(projected); } - ggml_tensor * future = f32(weights.future_tokens); - ggml_tensor * position_view = ggml_view_2d( - graph_context, weights.action_position, width, config.horizon, - weights.action_position->nb[1], 0); - ggml_tensor * position = f32(position_view); + ggml_tensor * future = f32(weights.future_tokens); + ggml_tensor * position_view = ggml_view_2d(graph_context, weights.action_position, width, config.horizon, + weights.action_position->nb[1], 0); + ggml_tensor * position = f32(position_view); // Qwen/projector inference and torch.randn run at BF16 in the released // script. The action head then enters CUDA autocast(float32). ggml_tensor * actions = bf16_roundtrip(noise_input); for (int step = 0; step < 4; ++step) { - ggml_tensor * timestep_projection = ggml_view_1d( - graph_context, timestep_projection_input, config.timestep_projection_dim, - static_cast(step) * config.timestep_projection_dim * sizeof(float)); - ggml_tensor * temb = linear(timestep_projection, weights.timestep_input_weight, - weights.timestep_input_bias); + ggml_tensor * timestep_projection = + ggml_view_1d(graph_context, timestep_projection_input, config.timestep_projection_dim, + static_cast(step) * config.timestep_projection_dim * sizeof(float)); + ggml_tensor * temb = + linear(timestep_projection, weights.timestep_input_weight, weights.timestep_input_bias); temb = ggml_silu(graph_context, temb); temb = linear(temb, weights.timestep_output_weight, weights.timestep_output_bias); - ggml_tensor * action_features = linear(actions, weights.action_input_weight, - weights.action_input_bias); - ggml_tensor * action_time = ggml_view_1d( - graph_context, action_time_input, width, - static_cast(step) * width * sizeof(float)); - action_time = ggml_repeat(graph_context, action_time, action_features); - action_features = ggml_concat(graph_context, action_features, action_time, 0); - action_features = linear(action_features, weights.action_time_mix_weight, - weights.action_time_mix_bias); + ggml_tensor * action_features = linear(actions, weights.action_input_weight, weights.action_input_bias); + ggml_tensor * action_time = ggml_view_1d(graph_context, action_time_input, width, + static_cast(step) * width * sizeof(float)); + action_time = ggml_repeat(graph_context, action_time, action_features); + action_features = ggml_concat(graph_context, action_features, action_time, 0); + action_features = linear(action_features, weights.action_time_mix_weight, weights.action_time_mix_bias); action_features = ggml_silu(graph_context, action_features); - action_features = linear(action_features, weights.action_output_weight, - weights.action_output_bias); + action_features = linear(action_features, weights.action_output_weight, weights.action_output_bias); action_features = ggml_add(graph_context, action_features, position); ggml_tensor * hidden = ggml_concat(graph_context, future, action_features, 1); for (int block_index = 0; block_index < config.block_count; ++block_index) { const PIV3BlockWeights & block = weights.blocks[static_cast(block_index)]; - ggml_tensor * normalized = ada_norm(hidden, temb, block); - ggml_tensor * attended = attention( - normalized, projected_hidden_states[static_cast(block_index)], block); - hidden = ggml_add(graph_context, hidden, attended); + ggml_tensor * normalized = ada_norm(hidden, temb, block); + ggml_tensor * attended = + attention(normalized, projected_hidden_states[static_cast(block_index)], block); + hidden = ggml_add(graph_context, hidden, attended); ggml_tensor * ff = ggml_norm(graph_context, hidden, config.ada_norm_epsilon); - ff = linear(ff, block.feed_forward_input_weight, block.feed_forward_input_bias); - ff = ggml_gelu(graph_context, ff); - ff = linear(ff, block.feed_forward_output_weight, block.feed_forward_output_bias); - hidden = ggml_add(graph_context, hidden, ff); + ff = linear(ff, block.feed_forward_input_weight, block.feed_forward_input_bias); + ff = ggml_gelu(graph_context, ff); + ff = linear(ff, block.feed_forward_output_weight, block.feed_forward_output_bias); + hidden = ggml_add(graph_context, hidden, ff); } // The released legacy sampler calls DiT with return_pre_output=true. // norm_out/proj_out_1/proj_out_2 are therefore intentionally inactive. - hidden = ggml_relu(graph_context, - linear(hidden, weights.velocity_input_weight, - weights.velocity_input_bias)); - hidden = linear(hidden, weights.velocity_output_weight, - weights.velocity_output_bias); - ggml_tensor * velocity = ggml_view_2d( - graph_context, hidden, config.action_dim, config.horizon, hidden->nb[1], - static_cast(config.future_token_count) * hidden->nb[1]); - actions = ggml_add(graph_context, actions, - ggml_scale(graph_context, velocity, config.euler_dt)); + hidden = + ggml_relu(graph_context, linear(hidden, weights.velocity_input_weight, weights.velocity_input_bias)); + hidden = linear(hidden, weights.velocity_output_weight, weights.velocity_output_bias); + ggml_tensor * velocity = + ggml_view_2d(graph_context, hidden, config.action_dim, config.horizon, hidden->nb[1], + static_cast(config.future_token_count) * hidden->nb[1]); + actions = ggml_add(graph_context, actions, ggml_scale(graph_context, velocity, config.euler_dt)); } output = actions; @@ -778,12 +585,12 @@ std::unique_ptr PIV3Policy::load(const std::string & path, int n_thr impl->verbosity = verbosity; try { backend_scheduler_config scheduler_config; - scheduler_config.max_nodes = static_cast(kGraphSize); - scheduler_config.parallel = false; + scheduler_config.max_nodes = static_cast(kGraphSize); + scheduler_config.parallel = false; scheduler_config.op_offload = true; backend_loader backend; - if (!backend.load(impl->backend_cpu, impl->backends, impl->scheduler, - impl->buft_policy, true, scheduler_config, verbosity)) { + if (!backend.load(impl->backend_cpu, impl->backends, impl->scheduler, impl->buft_policy, true, scheduler_config, + verbosity)) { error = "failed to initialize StarVLA PI_v3 backend: " + backend.error(); return nullptr; } @@ -798,16 +605,14 @@ std::unique_ptr PIV3Policy::load(const std::string & path, int n_thr error = "StarVLA PI_v3 policy GGUF has no tensors"; return nullptr; } - ggml_backend_buffer_set_usage(impl->loaded.model_buffer, - GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + ggml_backend_buffer_set_usage(impl->loaded.model_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); impl->timestep_table = timestep_projection_table(impl->config); - impl->action_table = action_time_table(impl->config); + impl->action_table = action_time_table(impl->config); if (verbosity >= 1) { - std::fprintf(stderr, - "%s: backend=%s qwen=%d width=%d layers=%d horizon=%d action_dim=%d profiles=%zu\n", - __func__, mode_name(impl->mode), impl->config.qwen_hidden_dim, - impl->config.dit_width, impl->config.block_count, impl->config.horizon, - impl->config.action_dim, impl->config.normalization.profiles.size()); + std::fprintf(stderr, "%s: backend=%s qwen=%d width=%d layers=%d horizon=%d action_dim=%d profiles=%zu\n", + __func__, backend_mode_name(impl->mode), impl->config.qwen_hidden_dim, impl->config.dit_width, + impl->config.block_count, impl->config.horizon, impl->config.action_dim, + impl->config.normalization.profiles.size()); } } catch (const std::exception & exception) { error = exception.what(); @@ -817,21 +622,16 @@ std::unique_ptr PIV3Policy::load(const std::string & path, int n_thr } bool PIV3Policy::evaluate(const float * qwen_hidden_states, size_t hidden_element_count, - const uint8_t * qwen_attention_mask, size_t mask_element_count, - const float * initial_noise, size_t noise_element_count, - std::vector & normalized_actions, std::string & error) { - return evaluate_internal(qwen_hidden_states, hidden_element_count, - qwen_attention_mask, mask_element_count, - initial_noise, noise_element_count, - normalized_actions, error); + const uint8_t * qwen_attention_mask, size_t mask_element_count, const float * initial_noise, + size_t noise_element_count, std::vector & normalized_actions, std::string & error) { + return evaluate_internal(qwen_hidden_states, hidden_element_count, qwen_attention_mask, mask_element_count, + initial_noise, noise_element_count, normalized_actions, error); } -bool PIV3Policy::evaluate_internal( - const float * qwen_hidden_states, size_t hidden_element_count, - const uint8_t * qwen_attention_mask, size_t mask_element_count, - const float * initial_noise, size_t noise_element_count, - std::vector & normalized_actions, - std::string & error) { +bool PIV3Policy::evaluate_internal(const float * qwen_hidden_states, size_t hidden_element_count, + const uint8_t * qwen_attention_mask, size_t mask_element_count, + const float * initial_noise, size_t noise_element_count, + std::vector & normalized_actions, std::string & error) { normalized_actions.clear(); error.clear(); if (impl_ == nullptr || impl_->scheduler == nullptr) { @@ -839,18 +639,16 @@ bool PIV3Policy::evaluate_internal( return false; } const size_t hidden_width = static_cast(impl_->config.qwen_hidden_dim); - const size_t layer_count = static_cast(impl_->config.qwen_layer_count); + const size_t layer_count = static_cast(impl_->config.qwen_layer_count); if (qwen_hidden_states == nullptr || qwen_attention_mask == nullptr || initial_noise == nullptr || - mask_element_count == 0 || - mask_element_count > static_cast(std::numeric_limits::max()) || + mask_element_count == 0 || mask_element_count > static_cast(std::numeric_limits::max()) || mask_element_count > std::numeric_limits::max() / hidden_width || mask_element_count * hidden_width > std::numeric_limits::max() / layer_count || hidden_element_count != mask_element_count * hidden_width * layer_count) { error = "StarVLA PI_v3 layerwise Qwen conditioning tensor or attention mask has an incompatible shape"; return false; } - const size_t expected_noise = - static_cast(impl_->config.horizon) * impl_->config.action_dim; + const size_t expected_noise = static_cast(impl_->config.horizon) * impl_->config.action_dim; if (noise_element_count != expected_noise) { error = "StarVLA PI_v3 initial-noise tensor has an incompatible shape"; return false; @@ -876,8 +674,7 @@ bool PIV3Policy::evaluate_internal( } try { - if (impl_->graph == nullptr || - impl_->conditioning_token_count != mask_element_count) { + if (impl_->graph == nullptr || impl_->conditioning_token_count != mask_element_count) { impl_->build_graph(mask_element_count); } } catch (const std::exception & exception) { @@ -885,25 +682,20 @@ bool PIV3Policy::evaluate_internal( return false; } - const int query_count = impl_->config.no_state_sequence_length; + const int query_count = impl_->config.no_state_sequence_length; const int padded_queries = GGML_PAD(query_count, kKQMaskPad); std::vector additive_mask(mask_element_count * static_cast(padded_queries), -std::numeric_limits::infinity()); for (int query = 0; query < query_count; ++query) { float * row = additive_mask.data() + static_cast(query) * mask_element_count; for (size_t token = 0; token < mask_element_count; ++token) { - row[token] = qwen_attention_mask[token] != 0 - ? 0.0f - : -std::numeric_limits::infinity(); + row[token] = qwen_attention_mask[token] != 0 ? 0.0f : -std::numeric_limits::infinity(); } } - ggml_backend_tensor_set(impl_->hidden_input, qwen_hidden_states, 0, - hidden_element_count * sizeof(float)); - ggml_backend_tensor_set(impl_->cross_mask_input, additive_mask.data(), 0, - additive_mask.size() * sizeof(float)); - ggml_backend_tensor_set(impl_->noise_input, initial_noise, 0, - noise_element_count * sizeof(float)); + ggml_backend_tensor_set(impl_->hidden_input, qwen_hidden_states, 0, hidden_element_count * sizeof(float)); + ggml_backend_tensor_set(impl_->cross_mask_input, additive_mask.data(), 0, additive_mask.size() * sizeof(float)); + ggml_backend_tensor_set(impl_->noise_input, initial_noise, 0, noise_element_count * sizeof(float)); ggml_backend_tensor_set(impl_->timestep_projection_input, impl_->timestep_table.data(), 0, impl_->timestep_table.size() * sizeof(float)); ggml_backend_tensor_set(impl_->action_time_input, impl_->action_table.data(), 0, @@ -917,8 +709,7 @@ bool PIV3Policy::evaluate_internal( } normalized_actions.resize(expected_noise); - ggml_backend_tensor_get(impl_->output, normalized_actions.data(), 0, - expected_noise * sizeof(float)); + ggml_backend_tensor_get(impl_->output, normalized_actions.data(), 0, expected_noise * sizeof(float)); if (std::any_of(normalized_actions.begin(), normalized_actions.end(), [](float value) { return !std::isfinite(value); })) { normalized_actions.clear(); @@ -928,17 +719,15 @@ bool PIV3Policy::evaluate_internal( return true; } -bool PIV3Policy::unnormalize(const std::vector & normalized_actions, - const std::string & profile_key_value, +bool PIV3Policy::unnormalize(const std::vector & normalized_actions, const std::string & profile_key_value, std::vector & actions, std::string & error) const { if (impl_ == nullptr) { actions.clear(); error = "StarVLA PI_v3 policy is not initialized"; return false; } - return denormalize_actions(impl_->config.normalization, profile_key_value, - normalized_actions, impl_->config.horizon, - impl_->config.action_dim, actions, error); + return denormalize_actions(impl_->config.normalization, profile_key_value, normalized_actions, + impl_->config.horizon, impl_->config.action_dim, actions, error); } const PIV3PolicyConfig & PIV3Policy::config() const { @@ -949,7 +738,7 @@ const PIV3PolicyConfig & PIV3Policy::config() const { } const char * PIV3Policy::backend_name() const { - return impl_ != nullptr ? mode_name(impl_->mode) : "unknown"; + return impl_ != nullptr ? backend_mode_name(impl_->mode) : "unknown"; } } // namespace robotcpp::starvla diff --git a/src/models/starvla/pi_v3_policy.h b/src/models/starvla/pi_v3_policy.h index 8177221..863e1fd 100644 --- a/src/models/starvla/pi_v3_policy.h +++ b/src/models/starvla/pi_v3_policy.h @@ -16,38 +16,38 @@ struct PIV3PolicyConfig { std::string text_filename; std::string mmproj_filename; - int qwen_hidden_dim = 0; + int qwen_hidden_dim = 0; int qwen_input_embedding_dim = 0; - int qwen_layer_count = 0; - int qwen_vocab_size = 0; + int qwen_layer_count = 0; + int qwen_vocab_size = 0; std::string cot_template; int image_count = 0; std::vector image_names; int image_processor_min_pixels = 0; int image_processor_max_pixels = 0; - int image_patch_size = 0; - int image_spatial_merge_size = 0; - int image_min_token_count = 0; - int image_max_token_count = 0; - - int dit_width = 0; - int block_count = 0; - int projector_count = 0; - int attention_head_count = 0; - int attention_head_dim = 0; - int feed_forward_dim = 0; - int mlp_hidden_dim = 0; - int action_dim = 0; - int horizon = 0; - int future_token_count = 0; - int action_position_count = 0; + int image_patch_size = 0; + int image_spatial_merge_size = 0; + int image_min_token_count = 0; + int image_max_token_count = 0; + + int dit_width = 0; + int block_count = 0; + int projector_count = 0; + int attention_head_count = 0; + int attention_head_dim = 0; + int feed_forward_dim = 0; + int mlp_hidden_dim = 0; + int action_dim = 0; + int horizon = 0; + int future_token_count = 0; + int action_position_count = 0; int no_state_sequence_length = 0; - int timestep_projection_dim = 0; - int num_timestep_buckets = 0; - int num_inference_timesteps = 0; - float ada_norm_epsilon = 0.0f; + int timestep_projection_dim = 0; + int num_timestep_buckets = 0; + int num_inference_timesteps = 0; + float ada_norm_epsilon = 0.0f; float projector_norm_epsilon = 0.0f; - float euler_dt = 0.0f; + float euler_dt = 0.0f; std::vector qwen_hidden_tuple_indices; std::vector timestep_ids; NormalizationConfig normalization; @@ -57,7 +57,7 @@ class PIV3Policy { public: ~PIV3Policy(); - PIV3Policy(const PIV3Policy &) = delete; + PIV3Policy(const PIV3Policy &) = delete; PIV3Policy & operator=(const PIV3Policy &) = delete; static std::unique_ptr load(const std::string & path, int n_threads, int verbosity, @@ -68,9 +68,8 @@ class PIV3Policy { // same full-chat token sequence. Non-zero mask entries participate in // every cross-attention block. The released checkpoint has no raw-state // runtime path. initial_noise is token-major [horizon, action_dim]. - bool evaluate(const float * qwen_hidden_states, size_t hidden_element_count, - const uint8_t * qwen_attention_mask, size_t mask_element_count, - const float * initial_noise, size_t noise_element_count, + bool evaluate(const float * qwen_hidden_states, size_t hidden_element_count, const uint8_t * qwen_attention_mask, + size_t mask_element_count, const float * initial_noise, size_t noise_element_count, std::vector & normalized_actions, std::string & error); bool unnormalize(const std::vector & normalized_actions, const std::string & profile_key, std::vector & actions, std::string & error) const; @@ -83,14 +82,9 @@ class PIV3Policy { explicit PIV3Policy(std::unique_ptr impl); - bool evaluate_internal(const float * qwen_hidden_states, - size_t hidden_element_count, - const uint8_t * qwen_attention_mask, - size_t mask_element_count, - const float * initial_noise, - size_t noise_element_count, - std::vector & normalized_actions, - std::string & error); + bool evaluate_internal(const float * qwen_hidden_states, size_t hidden_element_count, + const uint8_t * qwen_attention_mask, size_t mask_element_count, const float * initial_noise, + size_t noise_element_count, std::vector & normalized_actions, std::string & error); std::unique_ptr impl_; }; diff --git a/src/models/starvla/policy_gguf.h b/src/models/starvla/policy_gguf.h new file mode 100644 index 0000000..9de19d0 --- /dev/null +++ b/src/models/starvla/policy_gguf.h @@ -0,0 +1,144 @@ +#pragma once + +#include "ggml.h" +#include "gguf.h" +#include "models/starvla/normalization.h" + +#include +#include +#include +#include +#include +#include + +namespace robotcpp::starvla::detail { + +inline int require_key(gguf_context * gguf, const char * key, gguf_type type) { + const int index = gguf_find_key(gguf, key); + if (index < 0) { + throw std::runtime_error(std::string("missing required StarVLA GGUF metadata: ") + key); + } + if (gguf_get_kv_type(gguf, index) != type) { + throw std::runtime_error(std::string("invalid StarVLA GGUF metadata type: ") + key); + } + return index; +} + +inline std::string require_string(gguf_context * gguf, const char * key) { + return gguf_get_val_str(gguf, require_key(gguf, key, GGUF_TYPE_STRING)); +} + +inline int32_t require_i32(gguf_context * gguf, const char * key) { + return gguf_get_val_i32(gguf, require_key(gguf, key, GGUF_TYPE_INT32)); +} + +inline float require_f32(gguf_context * gguf, const char * key) { + return gguf_get_val_f32(gguf, require_key(gguf, key, GGUF_TYPE_FLOAT32)); +} + +inline bool require_bool(gguf_context * gguf, const char * key) { + return gguf_get_val_bool(gguf, require_key(gguf, key, GGUF_TYPE_BOOL)); +} + +inline int require_array(gguf_context * gguf, const char * key, gguf_type element_type) { + const int index = require_key(gguf, key, GGUF_TYPE_ARRAY); + if (gguf_get_arr_type(gguf, index) != element_type) { + throw std::runtime_error(std::string("invalid StarVLA GGUF array element type: ") + key); + } + return index; +} + +inline std::vector require_string_array(gguf_context * gguf, const char * key) { + const int index = require_array(gguf, key, GGUF_TYPE_STRING); + const size_t count = gguf_get_arr_n(gguf, index); + std::vector values; + values.reserve(count); + for (size_t i = 0; i < count; ++i) { + values.emplace_back(gguf_get_arr_str(gguf, index, i)); + } + return values; +} + +template +inline std::vector require_numeric_array(gguf_context * gguf, const char * key, gguf_type type) { + const int index = require_array(gguf, key, type); + const size_t count = gguf_get_arr_n(gguf, index); + if (count == 0) { + return {}; + } + const auto * data = static_cast(gguf_get_arr_data(gguf, index)); + if (data == nullptr) { + throw std::runtime_error(std::string("missing StarVLA GGUF array data: ") + key); + } + return std::vector(data, data + count); +} + +inline std::vector require_i32_array(gguf_context * gguf, const char * key) { + return require_numeric_array(gguf, key, GGUF_TYPE_INT32); +} + +inline std::vector require_f32_array(gguf_context * gguf, const char * key) { + return require_numeric_array(gguf, key, GGUF_TYPE_FLOAT32); +} + +inline std::vector require_bool_array(gguf_context * gguf, const char * key) { + const auto raw = require_numeric_array(gguf, key, GGUF_TYPE_BOOL); + std::vector values(raw.size()); + for (size_t i = 0; i < raw.size(); ++i) { + values[i] = raw[i] != 0 ? 1 : 0; + } + return values; +} + +inline std::string profile_key(int index, const char * suffix) { + return "starvla.normalization.profile." + std::to_string(index) + "." + suffix; +} + +inline NormalizationConfig require_normalization(gguf_context * gguf, int action_dim) { + NormalizationConfig config; + config.clip_actions = require_bool(gguf, "starvla.normalization.clip_actions"); + config.binary_threshold = require_f32(gguf, "starvla.normalization.binary_threshold"); + config.binary_comparison = require_string(gguf, "starvla.normalization.binary_comparison"); + config.continuous_dimensions = require_i32_array(gguf, "starvla.action.continuous_dimensions"); + config.binary_dimensions = require_i32_array(gguf, "starvla.action.binary_dimensions"); + + const int profile_count = require_i32(gguf, "starvla.normalization.profile_count"); + const auto keys = require_string_array(gguf, "starvla.normalization.profile_keys"); + if (profile_count <= 0 || keys.size() != static_cast(profile_count)) { + throw std::runtime_error("StarVLA normalization profile count is inconsistent"); + } + config.default_profile_key = keys.front(); + config.profiles.reserve(static_cast(profile_count)); + for (int index = 0; index < profile_count; ++index) { + NormalizationProfile profile; + profile.key = require_string(gguf, profile_key(index, "key").c_str()); + profile.action_q01 = require_f32_array(gguf, profile_key(index, "action_q01").c_str()); + profile.action_q99 = require_f32_array(gguf, profile_key(index, "action_q99").c_str()); + profile.action_mask = require_bool_array(gguf, profile_key(index, "action_mask").c_str()); + if (profile.key != keys[static_cast(index)]) { + throw std::runtime_error("StarVLA normalization profile order is inconsistent"); + } + config.profiles.push_back(std::move(profile)); + } + + std::string error; + if (!validate_normalization_config(config, action_dim, error)) { + throw std::runtime_error(error); + } + return config; +} + +inline bool has_shape(const ggml_tensor * tensor, std::initializer_list expected) { + if (tensor == nullptr || static_cast(ggml_n_dims(tensor)) != expected.size()) { + return false; + } + size_t dimension = 0; + for (const int64_t value : expected) { + if (tensor->ne[dimension++] != value) { + return false; + } + } + return true; +} + +} // namespace robotcpp::starvla::detail diff --git a/src/models/starvla/qwen3vl_bridge.cpp b/src/models/starvla/qwen3vl_bridge.cpp index 7a3735f..2c9bd12 100644 --- a/src/models/starvla/qwen3vl_bridge.cpp +++ b/src/models/starvla/qwen3vl_bridge.cpp @@ -1,5 +1,9 @@ #include "models/starvla/qwen3vl_bridge.h" +#ifdef ROBOTCPP_STARVLA_CUDA +#include "models/starvla/qwen_bf16_round_cuda.h" +#endif + #include "ggml.h" #include "gguf.h" #include "llama.h" @@ -20,27 +24,21 @@ namespace robotcpp::starvla { -bool qwen_vl_resolve_architecture(const std::string & text_architecture, - const std::string & projector_type, - QwenVLArchitecture & architecture, - std::string & error) { +bool qwen_vl_resolve_architecture(const std::string & text_architecture, const std::string & projector_type, + QwenVLArchitecture & architecture, std::string & error) { architecture = QwenVLArchitecture::unknown; error.clear(); - if (text_architecture == "qwen2vl" && - projector_type == "qwen2.5vl_merger") { + if (text_architecture == "qwen2vl" && projector_type == "qwen2.5vl_merger") { architecture = QwenVLArchitecture::qwen2_5_vl; return true; } - if (text_architecture == "qwen3vl" && - projector_type == "qwen3vl_merger") { + if (text_architecture == "qwen3vl" && projector_type == "qwen3vl_merger") { architecture = QwenVLArchitecture::qwen3_vl; return true; } - if (text_architecture != "qwen2vl" && - text_architecture != "qwen3vl") { + if (text_architecture != "qwen2vl" && text_architecture != "qwen3vl") { error = "unsupported Qwen-VL text architecture: " + text_architecture; - } else if (projector_type != "qwen2.5vl_merger" && - projector_type != "qwen3vl_merger") { + } else if (projector_type != "qwen2.5vl_merger" && projector_type != "qwen3vl_merger") { error = "unsupported Qwen-VL projector type: " + projector_type; } else { error = "Qwen-VL text and mmproj architectures do not match"; @@ -61,21 +59,14 @@ const char * qwen_vl_architecture_name(QwenVLArchitecture architecture) { } bool qwen_vl_is_final_norm_tensor_name(const char * name) noexcept { - return name != nullptr && - (std::strcmp(name, "result_norm") == 0 || - std::strcmp(name, "result_embd_pooled") == 0); + return name != nullptr && (std::strcmp(name, "result_norm") == 0 || std::strcmp(name, "result_embd_pooled") == 0); } -bool qwen_vl_hidden_state_source(QwenVLArchitecture architecture, - int decoder_layer_count, - int deepstack_layer_count, - int32_t hidden_tuple_index, - QwenVLHiddenStateSource & source, - std::string & error) { +bool qwen_vl_hidden_state_source(QwenVLArchitecture architecture, int decoder_layer_count, int deepstack_layer_count, + int32_t hidden_tuple_index, QwenVLHiddenStateSource & source, std::string & error) { source = QwenVLHiddenStateSource{}; error.clear(); - if (decoder_layer_count <= 0 || hidden_tuple_index <= 0 || - hidden_tuple_index > decoder_layer_count) { + if (decoder_layer_count <= 0 || hidden_tuple_index <= 0 || hidden_tuple_index > decoder_layer_count) { error = "Qwen-VL hidden-state tuple index is out of range"; return false; } @@ -85,23 +76,21 @@ bool qwen_vl_hidden_state_source(QwenVLArchitecture architecture, return false; } if (hidden_tuple_index == decoder_layer_count) { - source.kind = QwenVLHiddenStateSourceKind::final_norm; + source.kind = QwenVLHiddenStateSourceKind::final_norm; source.layer = -1; } else { - source.kind = QwenVLHiddenStateSourceKind::decoder_output; + source.kind = QwenVLHiddenStateSourceKind::decoder_output; source.layer = hidden_tuple_index - 1; } return true; } if (architecture == QwenVLArchitecture::qwen3_vl) { - if (deepstack_layer_count <= 0 || - deepstack_layer_count > decoder_layer_count) { + if (deepstack_layer_count <= 0 || deepstack_layer_count > decoder_layer_count) { error = "Qwen3-VL DeepStack layer count is incompatible with the model"; return false; } - source.kind = hidden_tuple_index <= deepstack_layer_count - ? QwenVLHiddenStateSourceKind::deepstack_output - : QwenVLHiddenStateSourceKind::decoder_output; + source.kind = hidden_tuple_index <= deepstack_layer_count ? QwenVLHiddenStateSourceKind::deepstack_output + : QwenVLHiddenStateSourceKind::decoder_output; source.layer = hidden_tuple_index - 1; return true; } @@ -109,14 +98,12 @@ bool qwen_vl_hidden_state_source(QwenVLArchitecture architecture, return false; } -bool qwen_vl_select_repetition_penalized_top1( - const float * logits, size_t vocab_size, - const std::vector & full_sequence, float repetition_penalty, - int32_t & token, std::string & error) { +bool qwen_vl_select_repetition_penalized_top1(const float * logits, size_t vocab_size, + const std::vector & full_sequence, float repetition_penalty, + int32_t & token, std::string & error) { token = -1; error.clear(); - if (logits == nullptr || vocab_size == 0 || - vocab_size > static_cast(INT32_MAX) || + if (logits == nullptr || vocab_size == 0 || vocab_size > static_cast(INT32_MAX) || !std::isfinite(repetition_penalty) || repetition_penalty <= 0.0f) { error = "Qwen-VL generation selector received an invalid contract"; return false; @@ -131,7 +118,7 @@ bool qwen_vl_select_repetition_penalized_top1( repeated[static_cast(value)] = 1; } - float best = -std::numeric_limits::infinity(); + float best = -std::numeric_limits::infinity(); int32_t best_token = -1; for (size_t index = 0; index < vocab_size; ++index) { float score = logits[index]; @@ -140,12 +127,11 @@ bool qwen_vl_select_repetition_penalized_top1( return false; } if (repeated[index] != 0) { - score = score < 0.0f ? score * repetition_penalty - : score / repetition_penalty; + score = score < 0.0f ? score * repetition_penalty : score / repetition_penalty; } // torch.argmax returns the first index on ties. if (best_token < 0 || score > best) { - best = score; + best = score; best_token = static_cast(index); } } @@ -160,14 +146,14 @@ bool qwen_vl_select_repetition_penalized_top1( namespace { void quiet_mtmd_log_callback(ggml_log_level level, const char * text, void * user_data) { - (void) user_data; + (void)user_data; if (level == GGML_LOG_LEVEL_ERROR) { std::fputs(text, stderr); } } struct PreparedMultimodalBatch { - size_t token_count = 0; + size_t token_count = 0; llama_pos position_count = 0; std::vector embeddings; std::vector positions; @@ -192,16 +178,16 @@ struct PreparedMultimodalBatch { struct BackendPlacement { bool accelerator_compute = false; - bool cpu_compute = false; + bool cpu_compute = false; }; struct LayerCapture { - BackendPlacement * placement = nullptr; - bool enabled = false; - bool bf16_residual_layer_boundaries = false; + BackendPlacement * placement = nullptr; + bool enabled = false; + bool bf16_residual_layer_boundaries = false; size_t expected_deepstack_layer_count = 0; - size_t token_count = 0; - size_t hidden_size = 0; + size_t token_count = 0; + size_t hidden_size = 0; std::vector layer_to_slot; std::vector deepstack_to_slot; int result_norm_slot = -1; @@ -209,10 +195,14 @@ struct LayerCapture { std::vector seen; std::vector rounded_layers; std::vector rounded_deepstack_layers; +#ifdef ROBOTCPP_STARVLA_CUDA + QwenBF16CaptureCuda cuda_capture; + size_t cuda_captured = 0; +#endif std::string error; void disable() { - enabled = false; + enabled = false; token_count = 0; hidden_size = 0; layer_to_slot.clear(); @@ -222,10 +212,28 @@ struct LayerCapture { seen.clear(); rounded_layers.clear(); rounded_deepstack_layers.clear(); +#ifdef ROBOTCPP_STARVLA_CUDA + cuda_captured = 0; +#endif error.clear(); } }; +bool finish_layer_capture(LayerCapture & capture) { +#ifdef ROBOTCPP_STARVLA_CUDA + if (capture.cuda_captured == 0) { + return true; + } + if (capture.cuda_captured != capture.seen.size()) { + capture.error = "Qwen-VL hidden states span multiple backends"; + return false; + } + return qwen_bf16_capture_download_cuda(capture.cuda_capture, capture.values.data(), capture.values.size(), + capture.error); +#endif + return true; +} + void begin_layer_boundary_tracking(LayerCapture & capture, size_t layer_count) { if (!capture.bf16_residual_layer_boundaries) { capture.rounded_layers.clear(); @@ -233,8 +241,7 @@ void begin_layer_boundary_tracking(LayerCapture & capture, size_t layer_count) { return; } capture.rounded_layers.assign(layer_count, uint8_t{0}); - capture.rounded_deepstack_layers.assign( - capture.expected_deepstack_layer_count, uint8_t{0}); + capture.rounded_deepstack_layers.assign(capture.expected_deepstack_layer_count, uint8_t{0}); } bool validate_layer_boundary_tracking(const LayerCapture & capture, std::string & error) { @@ -242,30 +249,25 @@ bool validate_layer_boundary_tracking(const LayerCapture & capture, std::string return true; } if (capture.rounded_layers.size() != capture.layer_to_slot.size() || - capture.rounded_deepstack_layers.size() != - capture.expected_deepstack_layer_count || + capture.rounded_deepstack_layers.size() != capture.expected_deepstack_layer_count || std::any_of(capture.rounded_layers.begin(), capture.rounded_layers.end(), [](uint8_t seen) { return seen != 1; }) || - std::any_of(capture.rounded_deepstack_layers.begin(), - capture.rounded_deepstack_layers.end(), + std::any_of(capture.rounded_deepstack_layers.begin(), capture.rounded_deepstack_layers.end(), [](uint8_t seen) { return seen != 1; })) { - error = - "Qwen3-VL BF16 residual-boundary roundtrip did not observe " - "every expected l_out/deepstack_out exactly once"; + error = "Qwen3-VL BF16 residual-boundary roundtrip did not observe " + "every expected l_out/deepstack_out exactly once"; return false; } return true; } bool observe_backend_placement(ggml_tensor * tensor, bool ask, void * user_data) { - if (!ask || tensor == nullptr || tensor->op == GGML_OP_NONE || tensor->buffer == nullptr || - user_data == nullptr) { + if (!ask || tensor == nullptr || tensor->op == GGML_OP_NONE || tensor->buffer == nullptr || user_data == nullptr) { return false; } - auto * placement = static_cast(user_data); + auto * placement = static_cast(user_data); ggml_backend_buffer_type_t buffer_type = ggml_backend_buffer_get_type(tensor->buffer); - ggml_backend_dev_t device = - buffer_type == nullptr ? nullptr : ggml_backend_buft_get_device(buffer_type); + ggml_backend_dev_t device = buffer_type == nullptr ? nullptr : ggml_backend_buft_get_device(buffer_type); if (device == nullptr) { return false; } @@ -293,8 +295,8 @@ int indexed_output_index(const char * name, const char * prefix) { if (*number == '\0') { return -1; } - errno = 0; - char * end = nullptr; + errno = 0; + char * end = nullptr; const long parsed = std::strtol(number, &end, 10); if (errno != 0 || end == number || *end != '\0' || parsed < 0 || parsed > INT_MAX) { return -1; @@ -312,16 +314,13 @@ bool observe_text_and_capture_layers(ggml_tensor * tensor, bool ask, void * user return false; } - int slot = -1; + int slot = -1; const int deepstack_layer = indexed_output_index(tensor->name, "deepstack_out-"); - const int layer = indexed_output_index(tensor->name, "l_out-"); - const bool is_result_norm = - qwen_vl_is_final_norm_tensor_name(tensor->name); + const int layer = indexed_output_index(tensor->name, "l_out-"); + const bool is_result_norm = qwen_vl_is_final_norm_tensor_name(tensor->name); const bool valid_deepstack_layer = - deepstack_layer >= 0 && - static_cast(deepstack_layer) < capture->deepstack_to_slot.size(); - const bool valid_layer = - layer >= 0 && static_cast(layer) < capture->layer_to_slot.size(); + deepstack_layer >= 0 && static_cast(deepstack_layer) < capture->deepstack_to_slot.size(); + const bool valid_layer = layer >= 0 && static_cast(layer) < capture->layer_to_slot.size(); if (valid_deepstack_layer) { slot = capture->deepstack_to_slot[static_cast(deepstack_layer)]; } else if (valid_layer) { @@ -329,12 +328,9 @@ bool observe_text_and_capture_layers(ggml_tensor * tensor, bool ask, void * user } else if (is_result_norm) { slot = capture->result_norm_slot; } - const bool round_layer = - capture->bf16_residual_layer_boundaries && valid_layer; - const bool round_deepstack = - capture->bf16_residual_layer_boundaries && valid_deepstack_layer && - static_cast(deepstack_layer) < - capture->expected_deepstack_layer_count; + const bool round_layer = capture->bf16_residual_layer_boundaries && valid_layer; + const bool round_deepstack = capture->bf16_residual_layer_boundaries && valid_deepstack_layer && + static_cast(deepstack_layer) < capture->expected_deepstack_layer_count; if (slot < 0 && !round_layer && !round_deepstack) { return false; } @@ -349,103 +345,125 @@ bool observe_text_and_capture_layers(ggml_tensor * tensor, bool ask, void * user return false; } if (capture->seen[static_cast(slot)] != 0) { - capture->error = - "Qwen3-VL emitted a requested hidden-state output more than once"; + capture->error = "Qwen3-VL emitted a requested hidden-state output more than once"; return false; } } - if (round_layer && - (static_cast(layer) >= capture->rounded_layers.size() || - capture->rounded_layers[static_cast(layer)] != 0)) { - capture->error = - "Qwen3-VL l_out BF16 roundtrip index is invalid or repeated"; + if (round_layer && (static_cast(layer) >= capture->rounded_layers.size() || + capture->rounded_layers[static_cast(layer)] != 0)) { + capture->error = "Qwen3-VL l_out BF16 roundtrip index is invalid or repeated"; return false; } - if (round_deepstack && - (static_cast(deepstack_layer) >= - capture->rounded_deepstack_layers.size() || - capture->rounded_deepstack_layers[static_cast(deepstack_layer)] != 0)) { - capture->error = - "Qwen3-VL deepstack_out BF16 roundtrip index is invalid or repeated"; + if (round_deepstack && (static_cast(deepstack_layer) >= capture->rounded_deepstack_layers.size() || + capture->rounded_deepstack_layers[static_cast(deepstack_layer)] != 0)) { + capture->error = "Qwen3-VL deepstack_out BF16 roundtrip index is invalid or repeated"; return false; } if (!ggml_is_contiguous(tensor) || tensor->ne[0] != static_cast(capture->hidden_size) || - tensor->ne[1] != static_cast(capture->token_count) || tensor->ne[2] != 1 || - tensor->ne[3] != 1) { + tensor->ne[1] != static_cast(capture->token_count) || tensor->ne[2] != 1 || tensor->ne[3] != 1) { capture->error = "Qwen3-VL hidden-state output has an incompatible shape or layout"; return false; } if (capture->hidden_size == 0 || - capture->token_count > - std::numeric_limits::max() / capture->hidden_size) { + capture->token_count > std::numeric_limits::max() / capture->hidden_size) { capture->error = "Qwen3-VL hidden-state capture size overflow"; return false; } const size_t count = capture->token_count * capture->hidden_size; if (count > std::numeric_limits::max() / sizeof(float) || - (slot >= 0 && - (count == 0 || static_cast(slot) >= capture->values.size() / count))) { + (slot >= 0 && (count == 0 || static_cast(slot) >= capture->values.size() / count))) { capture->error = "Qwen3-VL hidden-state capture byte range is invalid"; return false; } - std::vector rounded(count); - if (tensor->type == GGML_TYPE_F32) { + bool rounded_on_device = false; + bool captured_on_device = false; +#ifdef ROBOTCPP_STARVLA_CUDA + if ((round_layer || round_deepstack) && tensor->type == GGML_TYPE_F32) { + const QwenBF16RoundStatus status = qwen_bf16_round_cuda(tensor, count, capture->error); + if (status == QwenBF16RoundStatus::error) { + return false; + } + rounded_on_device = status == QwenBF16RoundStatus::success; + } + if (slot >= 0 && capture->seen.size() > 1 && tensor->type == GGML_TYPE_F32) { + const QwenBF16RoundStatus status = + qwen_bf16_capture_cuda(tensor, count, static_cast(slot) * count, capture->values.size(), + capture->cuda_capture, capture->error); + if (status == QwenBF16RoundStatus::error) { + return false; + } + captured_on_device = status == QwenBF16RoundStatus::success; + if (captured_on_device) { + ++capture->cuda_captured; + capture->seen[static_cast(slot)] = 1; + } + } +#endif + if (captured_on_device) { + if (round_layer) { + capture->rounded_layers[static_cast(layer)] = 1; + } else if (round_deepstack) { + capture->rounded_deepstack_layers[static_cast(deepstack_layer)] = 1; + } + return true; + } + std::vector rounded; + if (!rounded_on_device || slot >= 0) { + rounded.resize(count); + } + if (rounded_on_device) { + if (slot >= 0) { + ggml_backend_tensor_get(tensor, rounded.data(), 0, count * sizeof(float)); + } + } else if (tensor->type == GGML_TYPE_F32) { std::vector source(count); ggml_backend_tensor_get(tensor, source.data(), 0, count * sizeof(float)); for (size_t index = 0; index < count; ++index) { - rounded[index] = - ggml_bf16_to_fp32(ggml_fp32_to_bf16(source[index])); + rounded[index] = ggml_bf16_to_fp32(ggml_fp32_to_bf16(source[index])); } } else if (tensor->type == GGML_TYPE_F16) { if (round_layer || round_deepstack) { - capture->error = - "Qwen3-VL BF16 residual-boundary roundtrip requires F32 tensors"; + capture->error = "Qwen3-VL BF16 residual-boundary roundtrip requires F32 tensors"; return false; } std::vector source(count); - ggml_backend_tensor_get(tensor, source.data(), 0, - count * sizeof(ggml_fp16_t)); + ggml_backend_tensor_get(tensor, source.data(), 0, count * sizeof(ggml_fp16_t)); for (size_t index = 0; index < count; ++index) { - rounded[index] = ggml_bf16_to_fp32( - ggml_fp32_to_bf16(ggml_fp16_to_fp32(source[index]))); + rounded[index] = ggml_bf16_to_fp32(ggml_fp32_to_bf16(ggml_fp16_to_fp32(source[index]))); } } else if (tensor->type == GGML_TYPE_BF16) { if (round_layer || round_deepstack) { - capture->error = - "Qwen3-VL BF16 residual-boundary roundtrip requires F32 tensors"; + capture->error = "Qwen3-VL BF16 residual-boundary roundtrip requires F32 tensors"; return false; } std::vector source(count); - ggml_backend_tensor_get(tensor, source.data(), 0, - count * sizeof(ggml_bf16_t)); + ggml_backend_tensor_get(tensor, source.data(), 0, count * sizeof(ggml_bf16_t)); for (size_t index = 0; index < count; ++index) { rounded[index] = ggml_bf16_to_fp32(source[index]); } } else { - capture->error = std::string("unsupported Qwen3-VL hidden-state output type: ") + - ggml_type_name(tensor->type); + capture->error = + std::string("unsupported Qwen3-VL hidden-state output type: ") + ggml_type_name(tensor->type); return false; } + if ((round_layer || round_deepstack) && !rounded_on_device) { + ggml_backend_tensor_set(tensor, rounded.data(), 0, count * sizeof(float)); + } if (round_layer || round_deepstack) { - ggml_backend_tensor_set(tensor, rounded.data(), 0, - count * sizeof(float)); if (round_layer) { capture->rounded_layers[static_cast(layer)] = 1; } else { - capture->rounded_deepstack_layers[ - static_cast(deepstack_layer)] = 1; + capture->rounded_deepstack_layers[static_cast(deepstack_layer)] = 1; } } if (slot >= 0) { - float * destination = - capture->values.data() + static_cast(slot) * count; + float * destination = capture->values.data() + static_cast(slot) * count; std::copy(rounded.begin(), rounded.end(), destination); capture->seen[static_cast(slot)] = 1; } return true; } catch (const std::exception & exception) { - capture->error = std::string("failed to capture Qwen3-VL hidden-state output: ") + - exception.what(); + capture->error = std::string("failed to capture Qwen3-VL hidden-state output: ") + exception.what(); return false; } catch (...) { capture->error = "failed to capture Qwen3-VL hidden-state output"; @@ -463,7 +481,7 @@ int32_t decode_and_synchronize(llama_context * context, llama_batch batch) { } std::string model_metadata(const llama_model * model, const char * key) { - char value[256] = {}; + char value[256] = {}; const int32_t length = llama_model_meta_val_str(model, key, value, sizeof(value)); if (length < 0 || static_cast(length) >= sizeof(value)) { throw std::runtime_error(std::string("missing or oversized Qwen GGUF metadata: ") + key); @@ -473,8 +491,8 @@ std::string model_metadata(const llama_model * model, const char * key) { std::string gguf_string_metadata(const std::string & path, const char * key) { gguf_init_params params{}; - params.no_alloc = true; - params.ctx = nullptr; + params.no_alloc = true; + params.ctx = nullptr; gguf_context * gguf = gguf_init_from_file(path.c_str(), params); if (gguf == nullptr) { throw std::runtime_error("failed to read GGUF metadata: " + path); @@ -489,24 +507,22 @@ std::string gguf_string_metadata(const std::string & path, const char * key) { return value; } -std::vector tokenize(const llama_vocab * vocab, const std::string & text, - bool parse_special) { - const int32_t required = -llama_tokenize(vocab, text.data(), static_cast(text.size()), - nullptr, 0, false, parse_special); +std::vector tokenize(const llama_vocab * vocab, const std::string & text, bool parse_special) { + const int32_t required = + -llama_tokenize(vocab, text.data(), static_cast(text.size()), nullptr, 0, false, parse_special); if (required <= 0) { return {}; } std::vector tokens(static_cast(required)); - const int32_t count = llama_tokenize(vocab, text.data(), static_cast(text.size()), - tokens.data(), required, false, parse_special); + const int32_t count = llama_tokenize(vocab, text.data(), static_cast(text.size()), tokens.data(), required, + false, parse_special); if (count != required) { return {}; } return tokens; } -std::string apply_chat_template(const llama_model * model, - QwenVLArchitecture architecture, +std::string apply_chat_template(const llama_model * model, QwenVLArchitecture architecture, const std::string & content) { const char * chat_template = llama_model_chat_template(model, nullptr); if (chat_template == nullptr || chat_template[0] == '\0') { @@ -516,13 +532,10 @@ std::string apply_chat_template(const llama_model * model, {"system", "You are a helpful assistant."}, {"user", content.c_str()}, }; - const size_t message_offset = - architecture == QwenVLArchitecture::qwen2_5_vl ? 0U : 1U; - const size_t message_count = - architecture == QwenVLArchitecture::qwen2_5_vl ? 2U : 1U; + const size_t message_offset = architecture == QwenVLArchitecture::qwen2_5_vl ? 0U : 1U; + const size_t message_count = architecture == QwenVLArchitecture::qwen2_5_vl ? 2U : 1U; const int32_t required = - llama_chat_apply_template(chat_template, messages + message_offset, - message_count, true, nullptr, 0); + llama_chat_apply_template(chat_template, messages + message_offset, message_count, true, nullptr, 0); if (required < 0) { throw std::runtime_error("failed to size the Qwen chat-template output"); } @@ -530,9 +543,8 @@ std::string apply_chat_template(const llama_model * model, if (buffer.size() > static_cast(INT32_MAX)) { throw std::runtime_error("Qwen chat-template output is too large"); } - const int32_t written = llama_chat_apply_template( - chat_template, messages + message_offset, message_count, true, - buffer.data(), static_cast(buffer.size())); + const int32_t written = llama_chat_apply_template(chat_template, messages + message_offset, message_count, true, + buffer.data(), static_cast(buffer.size())); if (written != required) { throw std::runtime_error("failed to apply the Qwen chat template"); } @@ -540,21 +552,20 @@ std::string apply_chat_template(const llama_model * model, } struct PackedImageLayout { - size_t row_bytes = 0; + size_t row_bytes = 0; size_t stride_bytes = 0; size_t packed_bytes = 0; }; -bool validate_image(const Qwen3VLImageView & image, const Qwen3VLBridgeConfig & config, - PackedImageLayout & layout, std::string & error) { - (void) config; +bool validate_image(const Qwen3VLImageView & image, const Qwen3VLBridgeConfig & config, PackedImageLayout & layout, + std::string & error) { + (void)config; layout = PackedImageLayout{}; - if (image.data == nullptr || image.channels != 3 || image.width <= 0 || - image.height <= 0) { + if (image.data == nullptr || image.channels != 3 || image.width <= 0 || image.height <= 0) { error = "Qwen3-VL bridge requires a non-empty RGB image"; return false; } - const size_t width = static_cast(image.width); + const size_t width = static_cast(image.width); const size_t height = static_cast(image.height); if (width > std::numeric_limits::max() / 3U) { error = "Qwen3-VL image row size overflow"; @@ -562,19 +573,13 @@ bool validate_image(const Qwen3VLImageView & image, const Qwen3VLBridgeConfig & } layout.row_bytes = width * 3U; if (image.stride_bytes < 0 || - (image.stride_bytes > 0 && - static_cast(image.stride_bytes) < layout.row_bytes)) { + (image.stride_bytes > 0 && static_cast(image.stride_bytes) < layout.row_bytes)) { error = "Qwen3-VL image stride is smaller than a packed RGB row"; return false; } - layout.stride_bytes = image.stride_bytes > 0 - ? static_cast(image.stride_bytes) - : layout.row_bytes; + layout.stride_bytes = image.stride_bytes > 0 ? static_cast(image.stride_bytes) : layout.row_bytes; if (height > std::numeric_limits::max() / layout.row_bytes || - (height > 1U && - height - 1U > - (std::numeric_limits::max() - layout.row_bytes) / - layout.stride_bytes)) { + (height > 1U && height - 1U > (std::numeric_limits::max() - layout.row_bytes) / layout.stride_bytes)) { error = "Qwen3-VL image buffer size overflow"; return false; } @@ -582,22 +587,18 @@ bool validate_image(const Qwen3VLImageView & image, const Qwen3VLBridgeConfig & return true; } -std::vector pack_image(const Qwen3VLImageView & image, - const PackedImageLayout & layout) { +std::vector pack_image(const Qwen3VLImageView & image, const PackedImageLayout & layout) { std::vector packed(layout.packed_bytes); for (int row = 0; row < image.height; ++row) { std::memcpy(packed.data() + static_cast(row) * layout.row_bytes, - image.data + static_cast(row) * layout.stride_bytes, - layout.row_bytes); + image.data + static_cast(row) * layout.stride_bytes, layout.row_bytes); } return packed; } -void tokenize_multimodal_prompt(const Qwen3VLBridgeConfig & config, - QwenVLArchitecture architecture, +void tokenize_multimodal_prompt(const Qwen3VLBridgeConfig & config, QwenVLArchitecture architecture, const llama_model * model, mtmd_context * vision, - const std::vector & images, - const std::string & instruction, + const std::vector & images, const std::string & instruction, mtmd::input_chunks & chunks) { std::vector> packed_images; packed_images.reserve(images.size()); @@ -609,29 +610,26 @@ void tokenize_multimodal_prompt(const Qwen3VLBridgeConfig & config, throw std::runtime_error(validation_error); } packed_images.push_back(pack_image(image, layout)); - bitmaps.entries.emplace_back(static_cast(image.width), - static_cast(image.height), + bitmaps.entries.emplace_back(static_cast(image.width), static_cast(image.height), packed_images.back().data()); if (bitmaps.entries.back().ptr == nullptr) { throw std::runtime_error("failed to create a Qwen3-VL image bitmap"); } } - const std::string content = - build_qwen_media_content(images.size(), instruction, mtmd_default_marker()); - const std::string formatted = - apply_chat_template(model, architecture, content); + const std::string content = build_qwen_media_content(images.size(), instruction, mtmd_default_marker()); + const std::string formatted = apply_chat_template(model, architecture, content); mtmd_input_text input_text{}; - input_text.text = formatted.c_str(); - input_text.add_special = false; + input_text.text = formatted.c_str(); + input_text.add_special = false; input_text.parse_special = true; chunks.ptr.reset(mtmd_input_chunks_init()); if (chunks.ptr == nullptr) { throw std::runtime_error("failed to allocate Qwen3-VL multimodal input chunks"); } std::vector bitmap_ptrs = bitmaps.c_ptr(); - const int32_t tokenize_result = mtmd_tokenize( - vision, chunks.ptr.get(), &input_text, bitmap_ptrs.data(), bitmap_ptrs.size()); + const int32_t tokenize_result = + mtmd_tokenize(vision, chunks.ptr.get(), &input_text, bitmap_ptrs.data(), bitmap_ptrs.size()); if (tokenize_result != 0) { throw std::runtime_error("failed to tokenize the Qwen3-VL multimodal prompt"); } @@ -652,29 +650,26 @@ const char * compiled_backend_name() { struct Qwen3VLBridge::Impl { Qwen3VLBridgeConfig config; QwenVLArchitecture architecture = QwenVLArchitecture::unknown; - size_t deepstack_layer_count = 0; - llama_model * model = nullptr; - llama_context * context = nullptr; - mtmd_context * vision = nullptr; - const llama_vocab * vocab = nullptr; - bool backend_initialized = false; + size_t deepstack_layer_count = 0; + llama_model * model = nullptr; + llama_context * context = nullptr; + mtmd_context * vision = nullptr; + const llama_vocab * vocab = nullptr; + bool backend_initialized = false; BackendPlacement text_placement; BackendPlacement vision_placement; LayerCapture layer_capture; mutable std::string backend_name = "unknown"; void refresh_backend_name() const { - const bool accelerator = - text_placement.accelerator_compute && vision_placement.accelerator_compute; - const bool cpu = text_placement.cpu_compute || vision_placement.cpu_compute; + const bool accelerator = text_placement.accelerator_compute && vision_placement.accelerator_compute; + const bool cpu = text_placement.cpu_compute || vision_placement.cpu_compute; if (accelerator && !cpu) { backend_name = compiled_backend_name(); - } else if (!text_placement.accelerator_compute && - !vision_placement.accelerator_compute && + } else if (!text_placement.accelerator_compute && !vision_placement.accelerator_compute && text_placement.cpu_compute && vision_placement.cpu_compute) { backend_name = "cpu"; - } else if (text_placement.accelerator_compute || - vision_placement.accelerator_compute) { + } else if (text_placement.accelerator_compute || vision_placement.accelerator_compute) { backend_name = "mixed"; } else { backend_name = "unknown"; @@ -703,29 +698,25 @@ struct Qwen3VLBridge::Impl { namespace { -void copy_token_embedding(const ggml_tensor * token_embeddings, llama_token token, - size_t hidden_size, float * destination) { +void copy_token_embedding(const ggml_tensor * token_embeddings, llama_token token, size_t hidden_size, + float * destination) { if (token_embeddings == nullptr || destination == nullptr || token < 0 || - token_embeddings->ne[0] != static_cast(hidden_size) || - token >= token_embeddings->ne[1] || !ggml_is_contiguous(token_embeddings) || - token_embeddings->buffer == nullptr) { + token_embeddings->ne[0] != static_cast(hidden_size) || token >= token_embeddings->ne[1] || + !ggml_is_contiguous(token_embeddings) || token_embeddings->buffer == nullptr) { throw std::runtime_error("Qwen3-VL token embedding table is incompatible"); } const size_t row_stride = token_embeddings->nb[1]; - if (row_stride == 0 || - static_cast(token) > std::numeric_limits::max() / row_stride) { + if (row_stride == 0 || static_cast(token) > std::numeric_limits::max() / row_stride) { throw std::runtime_error("Qwen3-VL token embedding row offset overflow"); } const size_t row_offset = static_cast(token) * row_stride; switch (token_embeddings->type) { case GGML_TYPE_F32: - ggml_backend_tensor_get(token_embeddings, destination, row_offset, - hidden_size * sizeof(float)); + ggml_backend_tensor_get(token_embeddings, destination, row_offset, hidden_size * sizeof(float)); break; case GGML_TYPE_F16: { std::vector row(hidden_size); - ggml_backend_tensor_get(token_embeddings, row.data(), row_offset, - hidden_size * sizeof(ggml_fp16_t)); + ggml_backend_tensor_get(token_embeddings, row.data(), row_offset, hidden_size * sizeof(ggml_fp16_t)); for (size_t index = 0; index < hidden_size; ++index) { destination[index] = ggml_fp16_to_fp32(row[index]); } @@ -733,8 +724,7 @@ void copy_token_embedding(const ggml_tensor * token_embeddings, llama_token toke } case GGML_TYPE_BF16: { std::vector row(hidden_size); - ggml_backend_tensor_get(token_embeddings, row.data(), row_offset, - hidden_size * sizeof(ggml_bf16_t)); + ggml_backend_tensor_get(token_embeddings, row.data(), row_offset, hidden_size * sizeof(ggml_bf16_t)); for (size_t index = 0; index < hidden_size; ++index) { destination[index] = ggml_bf16_to_fp32(row[index]); } @@ -746,10 +736,8 @@ void copy_token_embedding(const ggml_tensor * token_embeddings, llama_token toke } } -PreparedMultimodalBatch prepare_multimodal_batch(const Qwen3VLBridgeConfig & config, - llama_model * model, - mtmd_context * vision, - const mtmd::input_chunks & chunks) { +PreparedMultimodalBatch prepare_multimodal_batch(const Qwen3VLBridgeConfig & config, llama_model * model, + mtmd_context * vision, const mtmd::input_chunks & chunks) { if (model == nullptr || vision == nullptr || !mtmd_decode_use_mrope(vision)) { throw std::runtime_error("Qwen3-VL single-batch decode requires M-RoPE components"); } @@ -757,9 +745,8 @@ PreparedMultimodalBatch prepare_multimodal_batch(const Qwen3VLBridgeConfig & con PreparedMultimodalBatch prepared; for (size_t chunk_index = 0; chunk_index < chunks.size(); ++chunk_index) { const mtmd_input_chunk * chunk = chunks[chunk_index]; - const size_t chunk_tokens = mtmd_input_chunk_get_n_tokens(chunk); - if (chunk_tokens == 0 || chunk_tokens > std::numeric_limits::max() - - prepared.token_count) { + const size_t chunk_tokens = mtmd_input_chunk_get_n_tokens(chunk); + if (chunk_tokens == 0 || chunk_tokens > std::numeric_limits::max() - prepared.token_count) { throw std::runtime_error("Qwen3-VL multimodal chunk has an invalid token count"); } prepared.token_count += chunk_tokens; @@ -769,10 +756,9 @@ PreparedMultimodalBatch prepare_multimodal_batch(const Qwen3VLBridgeConfig & con } const size_t hidden_size = static_cast(config.hidden_size); - const size_t input_size = static_cast(config.input_embedding_size); + const size_t input_size = static_cast(config.input_embedding_size); if (hidden_size == 0 || input_size != static_cast(llama_model_n_embd_inp(model)) || - input_size < hidden_size || - prepared.token_count > std::numeric_limits::max() / input_size) { + input_size < hidden_size || prepared.token_count > std::numeric_limits::max() / input_size) { throw std::runtime_error("Qwen3-VL input embedding dimensions are incompatible"); } if (prepared.token_count > std::numeric_limits::max() / 4U) { @@ -791,22 +777,20 @@ PreparedMultimodalBatch prepare_multimodal_batch(const Qwen3VLBridgeConfig & con } const ggml_tensor * token_embeddings = model->get_tensor("token_embd.weight"); - size_t token_offset = 0; - llama_pos position_offset = 0; + size_t token_offset = 0; + llama_pos position_offset = 0; for (size_t chunk_index = 0; chunk_index < chunks.size(); ++chunk_index) { - const mtmd_input_chunk * chunk = chunks[chunk_index]; - const size_t chunk_tokens = mtmd_input_chunk_get_n_tokens(chunk); + const mtmd_input_chunk * chunk = chunks[chunk_index]; + const size_t chunk_tokens = mtmd_input_chunk_get_n_tokens(chunk); const llama_pos chunk_positions = mtmd_input_chunk_get_n_pos(chunk); - if (chunk_positions <= 0 || - position_offset > std::numeric_limits::max() - chunk_positions) { + if (chunk_positions <= 0 || position_offset > std::numeric_limits::max() - chunk_positions) { throw std::runtime_error("Qwen3-VL multimodal positions overflow"); } const mtmd_input_chunk_type type = mtmd_input_chunk_get_type(chunk); if (type == MTMD_INPUT_CHUNK_TYPE_TEXT) { - size_t text_token_count = 0; - const llama_token * tokens = - mtmd_input_chunk_get_tokens_text(chunk, &text_token_count); + size_t text_token_count = 0; + const llama_token * tokens = mtmd_input_chunk_get_tokens_text(chunk, &text_token_count); if (tokens == nullptr || text_token_count != chunk_tokens || chunk_positions != static_cast(chunk_tokens)) { throw std::runtime_error("Qwen3-VL text chunk contract is incompatible"); @@ -816,18 +800,15 @@ PreparedMultimodalBatch prepare_multimodal_batch(const Qwen3VLBridgeConfig & con copy_token_embedding(token_embeddings, tokens[local_index], hidden_size, prepared.embeddings.data() + global_index * input_size); prepared.token_ids[global_index] = tokens[local_index]; - const llama_pos position = - position_offset + static_cast(local_index); + const llama_pos position = position_offset + static_cast(local_index); for (size_t axis = 0; axis < 3U; ++axis) { prepared.positions[axis * prepared.token_count + global_index] = position; } prepared.positions[prepared.token_count * 3U + global_index] = 0; } } else if (type == MTMD_INPUT_CHUNK_TYPE_IMAGE) { - const mtmd_image_tokens * image_tokens = - mtmd_input_chunk_get_tokens_image(chunk); - if (image_tokens == nullptr || - mtmd_image_tokens_get_n_tokens(image_tokens) != chunk_tokens) { + const mtmd_image_tokens * image_tokens = mtmd_input_chunk_get_tokens_image(chunk); + if (image_tokens == nullptr || mtmd_image_tokens_get_n_tokens(image_tokens) != chunk_tokens) { throw std::runtime_error("Qwen3-VL image chunk contract is incompatible"); } if (mtmd_encode_chunk(vision, chunk) != 0) { @@ -838,23 +819,18 @@ PreparedMultimodalBatch prepare_multimodal_batch(const Qwen3VLBridgeConfig & con throw std::runtime_error("Qwen3-VL image encoder returned no embeddings"); } const size_t image_element_count = chunk_tokens * input_size; - float * destination = - prepared.embeddings.data() + token_offset * input_size; + float * destination = prepared.embeddings.data() + token_offset * input_size; for (size_t element = 0; element < image_element_count; ++element) { - destination[element] = ggml_bf16_to_fp32( - ggml_fp32_to_bf16(image_embeddings[element])); + destination[element] = ggml_bf16_to_fp32(ggml_fp32_to_bf16(image_embeddings[element])); } for (size_t local_index = 0; local_index < chunk_tokens; ++local_index) { const size_t global_index = token_offset + local_index; - const mtmd_decoder_pos position = mtmd_image_tokens_get_decoder_pos( - image_tokens, position_offset, local_index); - prepared.positions[global_index] = static_cast(position.t); - prepared.positions[prepared.token_count + global_index] = - static_cast(position.y); - prepared.positions[prepared.token_count * 2U + global_index] = - static_cast(position.x); - prepared.positions[prepared.token_count * 3U + global_index] = - static_cast(position.z); + const mtmd_decoder_pos position = + mtmd_image_tokens_get_decoder_pos(image_tokens, position_offset, local_index); + prepared.positions[global_index] = static_cast(position.t); + prepared.positions[prepared.token_count + global_index] = static_cast(position.y); + prepared.positions[prepared.token_count * 2U + global_index] = static_cast(position.x); + prepared.positions[prepared.token_count * 3U + global_index] = static_cast(position.z); } } else { throw std::runtime_error("Qwen3-VL prompt contains an unsupported media chunk"); @@ -869,18 +845,13 @@ PreparedMultimodalBatch prepare_multimodal_batch(const Qwen3VLBridgeConfig & con return prepared; } -void export_prepared_inputs(const Qwen3VLBridgeConfig & config, - const llama_vocab * vocab, - const mtmd::input_chunks & chunks, - const PreparedMultimodalBatch & prepared, - std::vector & input_ids, - std::vector & attention_mask, +void export_prepared_inputs(const Qwen3VLBridgeConfig & config, const llama_vocab * vocab, + const mtmd::input_chunks & chunks, const PreparedMultimodalBatch & prepared, + std::vector & input_ids, std::vector & attention_mask, std::vector & image_grid_thw) { - const std::vector image_pad_tokens = - tokenize(vocab, "<|image_pad|>", true); + const std::vector image_pad_tokens = tokenize(vocab, "<|image_pad|>", true); if (image_pad_tokens.size() != 1) { - throw std::runtime_error( - "Qwen3-VL vocabulary does not expose a unique <|image_pad|> token"); + throw std::runtime_error("Qwen3-VL vocabulary does not expose a unique <|image_pad|> token"); } if (config.image_spatial_merge_size <= 0) { throw std::runtime_error("Qwen3-VL image spatial merge size is invalid"); @@ -888,8 +859,7 @@ void export_prepared_inputs(const Qwen3VLBridgeConfig & config, input_ids.reserve(prepared.token_ids.size()); for (llama_token token : prepared.token_ids) { - input_ids.push_back(token < 0 ? static_cast(image_pad_tokens.front()) - : static_cast(token)); + input_ids.push_back(token < 0 ? static_cast(image_pad_tokens.front()) : static_cast(token)); } attention_mask.assign(prepared.token_count, uint8_t{1}); @@ -900,30 +870,25 @@ void export_prepared_inputs(const Qwen3VLBridgeConfig & config, if (mtmd_input_chunk_get_type(chunk) != MTMD_INPUT_CHUNK_TYPE_IMAGE) { continue; } - const mtmd_image_tokens * image_tokens = - mtmd_input_chunk_get_tokens_image(chunk); + const mtmd_image_tokens * image_tokens = mtmd_input_chunk_get_tokens_image(chunk); if (image_tokens == nullptr) { throw std::runtime_error("Qwen3-VL image chunk has no token grid"); } - const size_t image_token_count = - mtmd_image_tokens_get_n_tokens(image_tokens); - uint32_t max_x = 0; - uint32_t max_y = 0; + const size_t image_token_count = mtmd_image_tokens_get_n_tokens(image_tokens); + uint32_t max_x = 0; + uint32_t max_y = 0; for (size_t token = 0; token < image_token_count; ++token) { - const mtmd_decoder_pos position = - mtmd_image_tokens_get_decoder_pos(image_tokens, 0, token); + const mtmd_decoder_pos position = mtmd_image_tokens_get_decoder_pos(image_tokens, 0, token); if (position.t != 0 || position.z != 0) { - throw std::runtime_error( - "Qwen3-VL image token grid does not use the expected M-RoPE layout"); + throw std::runtime_error("Qwen3-VL image token grid does not use the expected M-RoPE layout"); } max_x = std::max(max_x, position.x); max_y = std::max(max_y, position.y); } - const size_t merged_width = static_cast(max_x) + 1U; + const size_t merged_width = static_cast(max_x) + 1U; const size_t merged_height = static_cast(max_y) + 1U; - const size_t merge = static_cast(config.image_spatial_merge_size); - if (merged_width == 0 || merged_height == 0 || - merged_width > static_cast(INT64_MAX) / merge || + const size_t merge = static_cast(config.image_spatial_merge_size); + if (merged_width == 0 || merged_height == 0 || merged_width > static_cast(INT64_MAX) / merge || merged_height > static_cast(INT64_MAX) / merge || merged_width > std::numeric_limits::max() / merged_height || merged_width * merged_height != image_token_count) { @@ -945,18 +910,15 @@ Qwen3VLBridge::Qwen3VLBridge(std::unique_ptr impl) : impl_(std::move(impl) Qwen3VLBridge::~Qwen3VLBridge() = default; -std::unique_ptr Qwen3VLBridge::load(const Qwen3VLBridgeConfig & config, - std::string & error) { +std::unique_ptr Qwen3VLBridge::load(const Qwen3VLBridgeConfig & config, std::string & error) { error.clear(); const bool action_config_valid = config.action_token.empty() ? config.action_token_id == -1 : config.action_token_id >= 0; if (config.text_path.empty() || config.mmproj_path.empty() || config.bundle_uuid.empty() || - config.hidden_size <= 0 || config.input_embedding_size <= 0 || config.vocab_size <= 0 || - !action_config_valid || + config.hidden_size <= 0 || config.input_embedding_size <= 0 || config.vocab_size <= 0 || !action_config_valid || config.expected_image_count <= 0 || config.image_min_tokens <= 0 || - config.image_max_tokens < config.image_min_tokens || - config.image_spatial_merge_size <= 0 || config.n_ctx <= 0 || - config.n_batch <= 0) { + config.image_max_tokens < config.image_min_tokens || config.image_spatial_merge_size <= 0 || + config.n_ctx <= 0 || config.n_batch <= 0) { error = "Qwen3-VL bridge configuration is incomplete"; return nullptr; } @@ -964,19 +926,17 @@ std::unique_ptr Qwen3VLBridge::load(const Qwen3VLBridgeConfig & c std::unique_ptr impl(new Impl()); impl->config = config; try { - const std::string mmproj_uuid = - gguf_string_metadata(config.mmproj_path, "general.source.uuid"); + const std::string mmproj_uuid = gguf_string_metadata(config.mmproj_path, "general.source.uuid"); if (mmproj_uuid != config.bundle_uuid) { throw std::runtime_error("Qwen3-VL mmproj bundle UUID does not match the policy"); } - const std::string projector_type = - gguf_string_metadata(config.mmproj_path, "clip.projector_type"); + const std::string projector_type = gguf_string_metadata(config.mmproj_path, "clip.projector_type"); llama_backend_init(); - impl->backend_initialized = true; + impl->backend_initialized = true; llama_model_params model_params = llama_model_default_params(); - model_params.n_gpu_layers = -1; - impl->model = llama_model_load_from_file(config.text_path.c_str(), model_params); + model_params.n_gpu_layers = -1; + impl->model = llama_model_load_from_file(config.text_path.c_str(), model_params); if (impl->model == nullptr) { throw std::runtime_error("failed to load Qwen3-VL text GGUF: " + config.text_path); } @@ -984,9 +944,8 @@ std::unique_ptr Qwen3VLBridge::load(const Qwen3VLBridgeConfig & c throw std::runtime_error("Qwen3-VL text bundle UUID does not match the policy"); } std::string profile_error; - if (!qwen_vl_resolve_architecture( - model_metadata(impl->model, "general.architecture"), - projector_type, impl->architecture, profile_error)) { + if (!qwen_vl_resolve_architecture(model_metadata(impl->model, "general.architecture"), projector_type, + impl->architecture, profile_error)) { throw std::runtime_error(profile_error); } if (llama_model_n_embd_out(impl->model) != config.hidden_size || @@ -994,21 +953,15 @@ std::unique_ptr Qwen3VLBridge::load(const Qwen3VLBridgeConfig & c throw std::runtime_error("Qwen3-VL text embedding dimensions do not match the policy"); } if (config.input_embedding_size % config.hidden_size != 0) { - throw std::runtime_error( - "Qwen-VL input embedding width is not an integral hidden-state layout"); + throw std::runtime_error("Qwen-VL input embedding width is not an integral hidden-state layout"); } - const int deepstack_layer_count = - config.input_embedding_size / config.hidden_size - 1; - if ((impl->architecture == QwenVLArchitecture::qwen2_5_vl && - deepstack_layer_count != 0) || + const int deepstack_layer_count = config.input_embedding_size / config.hidden_size - 1; + if ((impl->architecture == QwenVLArchitecture::qwen2_5_vl && deepstack_layer_count != 0) || (impl->architecture == QwenVLArchitecture::qwen3_vl && - (deepstack_layer_count <= 0 || - deepstack_layer_count > llama_model_n_layer(impl->model)))) { - throw std::runtime_error( - "Qwen-VL input embedding layout does not match the detected architecture"); + (deepstack_layer_count <= 0 || deepstack_layer_count > llama_model_n_layer(impl->model)))) { + throw std::runtime_error("Qwen-VL input embedding layout does not match the detected architecture"); } - impl->deepstack_layer_count = - static_cast(deepstack_layer_count); + impl->deepstack_layer_count = static_cast(deepstack_layer_count); for (int layer = 0; layer < llama_model_n_layer(impl->model); ++layer) { ggml_backend_dev_t device = impl->model->dev_layer(layer); if (device == nullptr) { @@ -1017,8 +970,7 @@ std::unique_ptr Qwen3VLBridge::load(const Qwen3VLBridgeConfig & c const enum ggml_backend_dev_type type = ggml_backend_dev_type(device); if (type == GGML_BACKEND_DEVICE_TYPE_CPU) { impl->text_placement.cpu_compute = true; - } else if (type == GGML_BACKEND_DEVICE_TYPE_GPU || - type == GGML_BACKEND_DEVICE_TYPE_IGPU || + } else if (type == GGML_BACKEND_DEVICE_TYPE_GPU || type == GGML_BACKEND_DEVICE_TYPE_IGPU || type == GGML_BACKEND_DEVICE_TYPE_ACCEL) { impl->text_placement.accelerator_compute = true; } @@ -1027,8 +979,7 @@ std::unique_ptr Qwen3VLBridge::load(const Qwen3VLBridgeConfig & c const enum ggml_backend_dev_type type = ggml_backend_dev_type(output_device); if (type == GGML_BACKEND_DEVICE_TYPE_CPU) { impl->text_placement.cpu_compute = true; - } else if (type == GGML_BACKEND_DEVICE_TYPE_GPU || - type == GGML_BACKEND_DEVICE_TYPE_IGPU || + } else if (type == GGML_BACKEND_DEVICE_TYPE_GPU || type == GGML_BACKEND_DEVICE_TYPE_IGPU || type == GGML_BACKEND_DEVICE_TYPE_ACCEL) { impl->text_placement.accelerator_compute = true; } @@ -1042,39 +993,34 @@ std::unique_ptr Qwen3VLBridge::load(const Qwen3VLBridgeConfig & c throw std::runtime_error("Qwen3-VL text vocabulary size does not match the policy"); } if (!config.action_token.empty()) { - const std::vector action_tokens = - tokenize(impl->vocab, config.action_token, true); + const std::vector action_tokens = tokenize(impl->vocab, config.action_token, true); if (action_tokens.size() != 1 || action_tokens.front() != config.action_token_id) { - throw std::runtime_error( - "Qwen3-VL action token mapping does not match the policy metadata"); + throw std::runtime_error("Qwen3-VL action token mapping does not match the policy metadata"); } } llama_context_params context_params = llama_context_default_params(); - context_params.n_ctx = static_cast(config.n_ctx); - context_params.n_batch = static_cast(config.n_batch); + context_params.n_ctx = static_cast(config.n_ctx); + context_params.n_batch = static_cast(config.n_batch); // Layer capture expects one complete l_out/deepstack_out tensor per decode call. - context_params.n_ubatch = static_cast(config.n_batch); - context_params.n_threads = config.n_threads; + context_params.n_ubatch = static_cast(config.n_batch); + context_params.n_threads = config.n_threads; context_params.n_threads_batch = config.n_threads; - context_params.pooling_type = LLAMA_POOLING_TYPE_NONE; - context_params.embeddings = false; + context_params.pooling_type = LLAMA_POOLING_TYPE_NONE; + context_params.embeddings = false; // Match the official Qwen3-VL BF16 inference cache instead of llama's F16 default. context_params.type_k = GGML_TYPE_BF16; context_params.type_v = GGML_TYPE_BF16; - context_params.flash_attn_type = config.flash_text_attention - ? LLAMA_FLASH_ATTN_TYPE_ENABLED - : LLAMA_FLASH_ATTN_TYPE_DISABLED; - impl->layer_capture.placement = &impl->text_placement; - impl->layer_capture.bf16_residual_layer_boundaries = - config.bf16_residual_layer_boundaries; + context_params.flash_attn_type = + config.flash_text_attention ? LLAMA_FLASH_ATTN_TYPE_ENABLED : LLAMA_FLASH_ATTN_TYPE_DISABLED; + impl->layer_capture.placement = &impl->text_placement; + impl->layer_capture.bf16_residual_layer_boundaries = config.bf16_residual_layer_boundaries; if (config.bf16_residual_layer_boundaries) { - impl->layer_capture.expected_deepstack_layer_count = - impl->deepstack_layer_count; + impl->layer_capture.expected_deepstack_layer_count = impl->deepstack_layer_count; } - context_params.cb_eval = observe_text_and_capture_layers; + context_params.cb_eval = observe_text_and_capture_layers; context_params.cb_eval_user_data = &impl->layer_capture; - impl->context = llama_init_from_model(impl->model, context_params); + impl->context = llama_init_from_model(impl->model, context_params); if (impl->context == nullptr) { throw std::runtime_error("failed to create Qwen3-VL text context"); } @@ -1083,18 +1029,17 @@ std::unique_ptr Qwen3VLBridge::load(const Qwen3VLBridgeConfig & c } if (llama_n_batch(impl->context) != static_cast(config.n_batch) || llama_n_ubatch(impl->context) != static_cast(config.n_batch)) { - throw std::runtime_error( - "Qwen3-VL text context did not preserve the requested batch/ubatch contract"); + throw std::runtime_error("Qwen3-VL text context did not preserve the requested batch/ubatch contract"); } mtmd_context_params vision_params = mtmd_context_params_default(); - vision_params.use_gpu = true; - vision_params.print_timings = config.verbosity >= 1; - vision_params.n_threads = config.n_threads; - vision_params.image_min_tokens = config.image_min_tokens; - vision_params.image_max_tokens = config.image_max_tokens; - vision_params.cb_eval = observe_backend_placement; - vision_params.cb_eval_user_data = &impl->vision_placement; + vision_params.use_gpu = true; + vision_params.print_timings = config.verbosity >= 1; + vision_params.n_threads = config.n_threads; + vision_params.image_min_tokens = config.image_min_tokens; + vision_params.image_max_tokens = config.image_max_tokens; + vision_params.cb_eval = observe_backend_placement; + vision_params.cb_eval_user_data = &impl->vision_placement; mtmd_log_set(config.verbosity >= 1 ? nullptr : quiet_mtmd_log_callback, nullptr); impl->vision = mtmd_init_from_file(config.mmproj_path.c_str(), impl->model, vision_params); if (impl->vision == nullptr) { @@ -1112,15 +1057,11 @@ std::unique_ptr Qwen3VLBridge::load(const Qwen3VLBridgeConfig & c "n_ctx=%u n_batch=%u n_ubatch=%u kv=bf16 " "text_native_graph_disable_requested=%s " "vision_native_graph_disable_requested=%s\n", - __func__, qwen_vl_architecture_name(impl->architecture), - impl->backend_name.c_str(), - llama_model_n_embd_out(impl->model), - llama_model_n_embd_inp(impl->model), - impl->deepstack_layer_count, config.expected_image_count, - config.image_min_tokens, config.image_max_tokens, - llama_n_ctx(impl->context), llama_n_batch(impl->context), - llama_n_ubatch(impl->context), - config.disable_text_backend_native_graphs ? "true" : "false", + __func__, qwen_vl_architecture_name(impl->architecture), impl->backend_name.c_str(), + llama_model_n_embd_out(impl->model), llama_model_n_embd_inp(impl->model), + impl->deepstack_layer_count, config.expected_image_count, config.image_min_tokens, + config.image_max_tokens, llama_n_ctx(impl->context), llama_n_batch(impl->context), + llama_n_ubatch(impl->context), config.disable_text_backend_native_graphs ? "true" : "false", config.disable_vision_backend_native_graphs ? "true" : "false"); } } catch (const std::exception & exception) { @@ -1131,13 +1072,12 @@ std::unique_ptr Qwen3VLBridge::load(const Qwen3VLBridgeConfig & c } bool Qwen3VLBridge::extract_token_embeddings(const std::vector & images, - const std::string & instruction, int32_t token_id, - size_t token_count, std::vector & embeddings, - std::string & error) { + const std::string & instruction, int32_t token_id, size_t token_count, + std::vector & embeddings, std::string & error) { embeddings.clear(); error.clear(); - if (impl_ == nullptr || impl_->model == nullptr || impl_->context == nullptr || - impl_->vision == nullptr || impl_->vocab == nullptr) { + if (impl_ == nullptr || impl_->model == nullptr || impl_->context == nullptr || impl_->vision == nullptr || + impl_->vocab == nullptr) { error = "Qwen3-VL bridge is not initialized"; return false; } @@ -1160,11 +1100,9 @@ bool Qwen3VLBridge::extract_token_embeddings(const std::vector try { mtmd::input_chunks chunks; - tokenize_multimodal_prompt(impl_->config, impl_->architecture, - impl_->model, impl_->vision, images, - instruction, chunks); - PreparedMultimodalBatch prepared = - prepare_multimodal_batch(impl_->config, impl_->model, impl_->vision, chunks); + tokenize_multimodal_prompt(impl_->config, impl_->architecture, impl_->model, impl_->vision, images, instruction, + chunks); + PreparedMultimodalBatch prepared = prepare_multimodal_batch(impl_->config, impl_->model, impl_->vision, chunks); std::vector matches; for (size_t index = 0; index < prepared.token_ids.size(); ++index) { @@ -1190,9 +1128,9 @@ bool Qwen3VLBridge::extract_token_embeddings(const std::vector throw std::runtime_error("Qwen3-VL model has no decoder layers"); } LayerCapture & capture = impl_->layer_capture; - capture.enabled = false; - capture.token_count = prepared.token_count; - capture.hidden_size = static_cast(impl_->config.hidden_size); + capture.enabled = false; + capture.token_count = prepared.token_count; + capture.hidden_size = static_cast(impl_->config.hidden_size); capture.layer_to_slot.assign(static_cast(layer_count), -1); capture.deepstack_to_slot.assign(static_cast(layer_count), -1); capture.result_norm_slot = -1; @@ -1210,9 +1148,9 @@ bool Qwen3VLBridge::extract_token_embeddings(const std::vector llama_memory_clear(llama_get_memory(impl_->context), true); llama_set_embeddings(impl_->context, true); - llama_batch batch = prepared.view(); + llama_batch batch = prepared.view(); const int32_t decode_result = decode_and_synchronize(impl_->context, batch); - capture.enabled = false; + capture.enabled = false; if (decode_result != 0) { throw std::runtime_error("failed to evaluate the Qwen3-VL multimodal batch"); } @@ -1224,16 +1162,13 @@ bool Qwen3VLBridge::extract_token_embeddings(const std::vector throw std::runtime_error(boundary_error); } if (capture.seen.size() != 1 || capture.seen.front() == 0) { - throw std::runtime_error( - "Qwen-VL did not expose the final conditioning output"); + throw std::runtime_error("Qwen-VL did not expose the final conditioning output"); } embeddings.resize(token_count * capture.hidden_size); for (size_t output_index = 0; output_index < matches.size(); ++output_index) { - const float * hidden = - capture.values.data() + matches[output_index] * capture.hidden_size; - std::copy_n(hidden, capture.hidden_size, - embeddings.data() + output_index * capture.hidden_size); + const float * hidden = capture.values.data() + matches[output_index] * capture.hidden_size; + std::copy_n(hidden, capture.hidden_size, embeddings.data() + output_index * capture.hidden_size); } impl_->layer_capture.disable(); llama_set_embeddings(impl_->context, false); @@ -1250,15 +1185,14 @@ bool Qwen3VLBridge::extract_token_embeddings(const std::vector } } -bool Qwen3VLBridge::extract_full_hidden_states( - const std::vector & images, const std::string & instruction, - std::vector & hidden_states, std::vector & attention_mask, - std::string & error) { +bool Qwen3VLBridge::extract_full_hidden_states(const std::vector & images, + const std::string & instruction, std::vector & hidden_states, + std::vector & attention_mask, std::string & error) { hidden_states.clear(); attention_mask.clear(); error.clear(); - if (impl_ == nullptr || impl_->model == nullptr || impl_->context == nullptr || - impl_->vision == nullptr || impl_->vocab == nullptr) { + if (impl_ == nullptr || impl_->model == nullptr || impl_->context == nullptr || impl_->vision == nullptr || + impl_->vocab == nullptr) { error = "Qwen3-VL bridge is not initialized"; return false; } @@ -1269,11 +1203,9 @@ bool Qwen3VLBridge::extract_full_hidden_states( try { mtmd::input_chunks chunks; - tokenize_multimodal_prompt(impl_->config, impl_->architecture, - impl_->model, impl_->vision, images, - instruction, chunks); - PreparedMultimodalBatch prepared = - prepare_multimodal_batch(impl_->config, impl_->model, impl_->vision, chunks); + tokenize_multimodal_prompt(impl_->config, impl_->architecture, impl_->model, impl_->vision, images, instruction, + chunks); + PreparedMultimodalBatch prepared = prepare_multimodal_batch(impl_->config, impl_->model, impl_->vision, chunks); if (prepared.token_count > llama_n_batch(impl_->context)) { throw std::runtime_error( "Qwen3-VL multimodal prompt exceeds n_batch; increase --n-batch for single-batch decode"); @@ -1292,9 +1224,9 @@ bool Qwen3VLBridge::extract_full_hidden_states( throw std::runtime_error("Qwen3-VL model has no decoder layers"); } LayerCapture & capture = impl_->layer_capture; - capture.enabled = false; - capture.token_count = prepared.token_count; - capture.hidden_size = hidden_size; + capture.enabled = false; + capture.token_count = prepared.token_count; + capture.hidden_size = hidden_size; capture.layer_to_slot.assign(static_cast(layer_count), -1); capture.deepstack_to_slot.assign(static_cast(layer_count), -1); capture.result_norm_slot = -1; @@ -1312,9 +1244,9 @@ bool Qwen3VLBridge::extract_full_hidden_states( llama_memory_clear(llama_get_memory(impl_->context), true); llama_set_embeddings(impl_->context, true); - llama_batch batch = prepared.view(); + llama_batch batch = prepared.view(); const int32_t decode_result = decode_and_synchronize(impl_->context, batch); - capture.enabled = false; + capture.enabled = false; if (decode_result != 0) { throw std::runtime_error("failed to evaluate the Qwen3-VL multimodal batch"); } @@ -1326,8 +1258,7 @@ bool Qwen3VLBridge::extract_full_hidden_states( throw std::runtime_error(boundary_error); } if (capture.seen.size() != 1 || capture.seen.front() == 0) { - throw std::runtime_error( - "Qwen-VL did not expose the final conditioning output"); + throw std::runtime_error("Qwen-VL did not expose the final conditioning output"); } hidden_states = std::move(capture.values); @@ -1348,15 +1279,16 @@ bool Qwen3VLBridge::extract_full_hidden_states( } } -bool Qwen3VLBridge::extract_layer_hidden_states( - const std::vector & images, const std::string & instruction, - const std::vector & hidden_tuple_indices, std::vector & hidden_states, - std::vector & attention_mask, std::string & error) { +bool Qwen3VLBridge::extract_layer_hidden_states(const std::vector & images, + const std::string & instruction, + const std::vector & hidden_tuple_indices, + std::vector & hidden_states, + std::vector & attention_mask, std::string & error) { hidden_states.clear(); attention_mask.clear(); error.clear(); - if (impl_ == nullptr || impl_->model == nullptr || impl_->context == nullptr || - impl_->vision == nullptr || impl_->vocab == nullptr) { + if (impl_ == nullptr || impl_->model == nullptr || impl_->context == nullptr || impl_->vision == nullptr || + impl_->vocab == nullptr) { error = "Qwen3-VL bridge is not initialized"; return false; } @@ -1373,20 +1305,17 @@ bool Qwen3VLBridge::extract_layer_hidden_states( } std::vector layer_to_slot(static_cast(model_layer_count), -1); std::vector deepstack_to_slot(static_cast(model_layer_count), -1); - if (impl_->deepstack_layer_count > - static_cast(std::numeric_limits::max())) { + if (impl_->deepstack_layer_count > static_cast(std::numeric_limits::max())) { error = "Qwen-VL DeepStack layer count exceeds the supported range"; return false; } - const int deepstack_layer_count = - static_cast(impl_->deepstack_layer_count); - int result_norm_slot = -1; + const int deepstack_layer_count = static_cast(impl_->deepstack_layer_count); + int result_norm_slot = -1; for (size_t slot = 0; slot < hidden_tuple_indices.size(); ++slot) { const int32_t tuple_index = hidden_tuple_indices[slot]; QwenVLHiddenStateSource source; - if (!qwen_vl_hidden_state_source( - impl_->architecture, model_layer_count, deepstack_layer_count, - tuple_index, source, error)) { + if (!qwen_vl_hidden_state_source(impl_->architecture, model_layer_count, deepstack_layer_count, tuple_index, + source, error)) { return false; } if (source.kind == QwenVLHiddenStateSourceKind::final_norm) { @@ -1402,9 +1331,7 @@ bool Qwen3VLBridge::extract_layer_hidden_states( return false; } std::vector & target = - source.kind == QwenVLHiddenStateSourceKind::deepstack_output - ? deepstack_to_slot - : layer_to_slot; + source.kind == QwenVLHiddenStateSourceKind::deepstack_output ? deepstack_to_slot : layer_to_slot; if (target[static_cast(source.layer)] >= 0) { error = "Qwen-VL hidden-state tuple indices must be unique"; return false; @@ -1414,11 +1341,9 @@ bool Qwen3VLBridge::extract_layer_hidden_states( try { mtmd::input_chunks chunks; - tokenize_multimodal_prompt(impl_->config, impl_->architecture, - impl_->model, impl_->vision, images, - instruction, chunks); - PreparedMultimodalBatch prepared = - prepare_multimodal_batch(impl_->config, impl_->model, impl_->vision, chunks); + tokenize_multimodal_prompt(impl_->config, impl_->architecture, impl_->model, impl_->vision, images, instruction, + chunks); + PreparedMultimodalBatch prepared = prepare_multimodal_batch(impl_->config, impl_->model, impl_->vision, chunks); if (prepared.token_count > llama_n_batch(impl_->context)) { throw std::runtime_error( "Qwen3-VL multimodal prompt exceeds n_batch; increase --n-batch for single-batch decode"); @@ -1428,47 +1353,46 @@ bool Qwen3VLBridge::extract_layer_hidden_states( throw std::runtime_error("Qwen3-VL multimodal prompt exceeds n_ctx"); } - const size_t hidden_size = static_cast(impl_->config.hidden_size); + const size_t hidden_size = static_cast(impl_->config.hidden_size); const size_t requested_layers = hidden_tuple_indices.size(); if (prepared.token_count > std::numeric_limits::max() / hidden_size || - prepared.token_count * hidden_size > - std::numeric_limits::max() / requested_layers) { + prepared.token_count * hidden_size > std::numeric_limits::max() / requested_layers) { throw std::runtime_error("Qwen3-VL layerwise hidden-state buffer size overflow"); } - LayerCapture & capture = impl_->layer_capture; - capture.enabled = false; - capture.token_count = prepared.token_count; - capture.hidden_size = hidden_size; - capture.layer_to_slot = layer_to_slot; + LayerCapture & capture = impl_->layer_capture; + capture.enabled = false; + capture.token_count = prepared.token_count; + capture.hidden_size = hidden_size; + capture.layer_to_slot = layer_to_slot; capture.deepstack_to_slot = deepstack_to_slot; - capture.result_norm_slot = result_norm_slot; + capture.result_norm_slot = result_norm_slot; capture.values.assign(requested_layers * prepared.token_count * hidden_size, 0.0f); capture.seen.assign(requested_layers, uint8_t{0}); - begin_layer_boundary_tracking(capture, - static_cast(model_layer_count)); + begin_layer_boundary_tracking(capture, static_cast(model_layer_count)); capture.error.clear(); capture.enabled = true; std::fill(prepared.outputs.begin(), prepared.outputs.end(), int8_t{1}); llama_memory_clear(llama_get_memory(impl_->context), true); llama_set_embeddings(impl_->context, true); - llama_batch batch = prepared.view(); + llama_batch batch = prepared.view(); const int32_t decode_result = decode_and_synchronize(impl_->context, batch); - capture.enabled = false; + capture.enabled = false; if (decode_result != 0) { throw std::runtime_error("failed to evaluate the Qwen3-VL multimodal batch"); } if (!capture.error.empty()) { throw std::runtime_error(capture.error); } + if (!finish_layer_capture(capture)) { + throw std::runtime_error(capture.error); + } std::string boundary_error; if (!validate_layer_boundary_tracking(capture, boundary_error)) { throw std::runtime_error(boundary_error); } - if (std::any_of(capture.seen.begin(), capture.seen.end(), - [](uint8_t seen) { return seen == 0; })) { - throw std::runtime_error( - "Qwen3-VL did not expose every requested hidden-state output"); + if (std::any_of(capture.seen.begin(), capture.seen.end(), [](uint8_t seen) { return seen == 0; })) { + throw std::runtime_error("Qwen3-VL did not expose every requested hidden-state output"); } hidden_states = std::move(capture.values); @@ -1489,15 +1413,13 @@ bool Qwen3VLBridge::extract_layer_hidden_states( } } -bool Qwen3VLBridge::generate_autoregressive( - const std::vector & images, - const std::string & instruction, - const QwenVLGenerationConfig & generation, - QwenVLGenerationResult & result, std::string & error) { +bool Qwen3VLBridge::generate_autoregressive(const std::vector & images, + const std::string & instruction, const QwenVLGenerationConfig & generation, + QwenVLGenerationResult & result, std::string & error) { result = QwenVLGenerationResult{}; error.clear(); - if (impl_ == nullptr || impl_->model == nullptr || impl_->context == nullptr || - impl_->vision == nullptr || impl_->vocab == nullptr) { + if (impl_ == nullptr || impl_->model == nullptr || impl_->context == nullptr || impl_->vision == nullptr || + impl_->vocab == nullptr) { error = "Qwen-VL bridge is not initialized"; return false; } @@ -1505,20 +1427,16 @@ bool Qwen3VLBridge::generate_autoregressive( error = "Qwen-VL image count does not match the policy"; return false; } - if (generation.max_length == 0 || - generation.max_length > static_cast(llama_n_ctx(impl_->context)) || - generation.max_length > static_cast(INT32_MAX) || - generation.top_k != 1 || generation.eos_token_ids.empty() || - !std::isfinite(generation.repetition_penalty) || + if (generation.max_length == 0 || generation.max_length > static_cast(llama_n_ctx(impl_->context)) || + generation.max_length > static_cast(INT32_MAX) || generation.top_k != 1 || + generation.eos_token_ids.empty() || !std::isfinite(generation.repetition_penalty) || generation.repetition_penalty <= 0.0f) { error = "Qwen-VL autoregressive generation configuration is incompatible"; return false; } - std::vector eos_seen(static_cast(impl_->config.vocab_size), - uint8_t{0}); + std::vector eos_seen(static_cast(impl_->config.vocab_size), uint8_t{0}); for (int32_t eos : generation.eos_token_ids) { - if (eos < 0 || eos >= impl_->config.vocab_size || - eos_seen[static_cast(eos)] != 0) { + if (eos < 0 || eos >= impl_->config.vocab_size || eos_seen[static_cast(eos)] != 0) { error = "Qwen-VL generation EOS token set is invalid"; return false; } @@ -1527,39 +1445,31 @@ bool Qwen3VLBridge::generate_autoregressive( try { mtmd::input_chunks chunks; - tokenize_multimodal_prompt(impl_->config, impl_->architecture, - impl_->model, impl_->vision, images, - instruction, chunks); - PreparedMultimodalBatch prepared = - prepare_multimodal_batch(impl_->config, impl_->model, impl_->vision, - chunks); + tokenize_multimodal_prompt(impl_->config, impl_->architecture, impl_->model, impl_->vision, images, instruction, + chunks); + PreparedMultimodalBatch prepared = prepare_multimodal_batch(impl_->config, impl_->model, impl_->vision, chunks); if (prepared.token_count > llama_n_batch(impl_->context)) { - throw std::runtime_error( - "Qwen-VL multimodal prompt exceeds n_batch; increase --n-batch"); + throw std::runtime_error("Qwen-VL multimodal prompt exceeds n_batch; increase --n-batch"); } if (prepared.token_count > generation.max_length || prepared.token_count > static_cast(llama_n_ctx(impl_->context)) || - prepared.position_count > - static_cast(llama_n_ctx(impl_->context))) { - throw std::runtime_error( - "Qwen-VL multimodal prompt exceeds the FAST max_length/n_ctx contract"); + prepared.position_count > static_cast(llama_n_ctx(impl_->context))) { + throw std::runtime_error("Qwen-VL multimodal prompt exceeds the FAST max_length/n_ctx contract"); } std::vector input_ids; std::vector attention_mask; std::vector image_grid_thw; - export_prepared_inputs(impl_->config, impl_->vocab, chunks, prepared, - input_ids, attention_mask, image_grid_thw); + export_prepared_inputs(impl_->config, impl_->vocab, chunks, prepared, input_ids, attention_mask, + image_grid_thw); if (input_ids.size() != prepared.token_count) { - throw std::runtime_error( - "Qwen-VL multimodal prompt token export is inconsistent"); + throw std::runtime_error("Qwen-VL multimodal prompt token export is inconsistent"); } result.prompt_token_count = prepared.token_count; result.full_sequence.reserve(generation.max_length); for (int64_t input_id : input_ids) { if (input_id < 0 || input_id >= impl_->config.vocab_size) { - throw std::runtime_error( - "Qwen-VL multimodal prompt contains an out-of-vocabulary token"); + throw std::runtime_error("Qwen-VL multimodal prompt contains an out-of-vocabulary token"); } result.full_sequence.push_back(static_cast(input_id)); } @@ -1572,56 +1482,43 @@ bool Qwen3VLBridge::generate_autoregressive( llama_set_embeddings(impl_->context, false); llama_memory_clear(llama_get_memory(impl_->context), true); std::fill(prepared.outputs.begin(), prepared.outputs.end(), int8_t{0}); - prepared.outputs.back() = 1; + prepared.outputs.back() = 1; llama_batch prompt_batch = prepared.view(); if (decode_and_synchronize(impl_->context, prompt_batch) != 0) { - throw std::runtime_error( - "failed to evaluate the Qwen-VL autoregressive prompt"); + throw std::runtime_error("failed to evaluate the Qwen-VL autoregressive prompt"); } while (result.full_sequence.size() < generation.max_length) { const float * logits = llama_get_logits_ith(impl_->context, -1); - int32_t next = -1; + int32_t next = -1; std::string selection_error; - if (!qwen_vl_select_repetition_penalized_top1( - logits, static_cast(impl_->config.vocab_size), - result.full_sequence, generation.repetition_penalty, next, - selection_error)) { + if (!qwen_vl_select_repetition_penalized_top1(logits, static_cast(impl_->config.vocab_size), + result.full_sequence, generation.repetition_penalty, next, + selection_error)) { throw std::runtime_error(selection_error); } result.full_sequence.push_back(next); result.continuation.push_back(next); - if (eos_seen[static_cast(next)] != 0 || - result.full_sequence.size() == generation.max_length) { + if (eos_seen[static_cast(next)] != 0 || result.full_sequence.size() == generation.max_length) { break; } const size_t generation_index = result.continuation.size() - 1U; if (generation_index > - static_cast(std::numeric_limits::max() - - prepared.position_count)) { - throw std::runtime_error( - "Qwen-VL autoregressive M-RoPE position overflow"); + static_cast(std::numeric_limits::max() - prepared.position_count)) { + throw std::runtime_error("Qwen-VL autoregressive M-RoPE position overflow"); } - llama_token token = static_cast(next); - llama_pos position = - prepared.position_count + static_cast(generation_index); - int32_t sequence_count = 1; + llama_token token = static_cast(next); + llama_pos position = prepared.position_count + static_cast(generation_index); + int32_t sequence_count = 1; llama_seq_id sequence_value = 0; - llama_seq_id * sequence = &sequence_value; - int8_t output = 1; + llama_seq_id * sequence = &sequence_value; + int8_t output = 1; llama_batch token_batch{ - 1, - &token, - nullptr, - &position, - &sequence_count, - &sequence, - &output, + 1, &token, nullptr, &position, &sequence_count, &sequence, &output, }; if (decode_and_synchronize(impl_->context, token_batch) != 0) { - throw std::runtime_error( - "failed to evaluate an incremental Qwen-VL generation token"); + throw std::runtime_error("failed to evaluate an incremental Qwen-VL generation token"); } } @@ -1634,7 +1531,7 @@ bool Qwen3VLBridge::generate_autoregressive( llama_set_embeddings(impl_->context, false); llama_memory_clear(llama_get_memory(impl_->context), true); result = QwenVLGenerationResult{}; - error = exception.what(); + error = exception.what(); return false; } } @@ -1656,10 +1553,6 @@ const char * Qwen3VLBridge::backend_name() const { return impl_->backend_name.c_str(); } -const char * Qwen3VLBridge::text_attention_mode_name() const { - return impl_ != nullptr && impl_->config.flash_text_attention ? "flash" : "non_flash"; -} - QwenVLArchitecture Qwen3VLBridge::architecture() const { return impl_ != nullptr ? impl_->architecture : QwenVLArchitecture::unknown; } diff --git a/src/models/starvla/qwen3vl_bridge.h b/src/models/starvla/qwen3vl_bridge.h index ec6f3eb..2c04e5e 100644 --- a/src/models/starvla/qwen3vl_bridge.h +++ b/src/models/starvla/qwen3vl_bridge.h @@ -21,18 +21,15 @@ enum class QwenVLHiddenStateSourceKind { }; struct QwenVLHiddenStateSource { - QwenVLHiddenStateSourceKind kind = - QwenVLHiddenStateSourceKind::decoder_output; - int layer = -1; + QwenVLHiddenStateSourceKind kind = QwenVLHiddenStateSourceKind::decoder_output; + int layer = -1; }; // Resolve the paired llama.cpp text and mtmd projector profiles. StarVLA // supports Qwen2.5-VL and Qwen3-VL only; mismatched text/mmproj files fail // before either component is evaluated. -bool qwen_vl_resolve_architecture(const std::string & text_architecture, - const std::string & projector_type, - QwenVLArchitecture & architecture, - std::string & error); +bool qwen_vl_resolve_architecture(const std::string & text_architecture, const std::string & projector_type, + QwenVLArchitecture & architecture, std::string & error); const char * qwen_vl_architecture_name(QwenVLArchitecture architecture); @@ -44,42 +41,37 @@ bool qwen_vl_is_final_norm_tensor_name(const char * name) noexcept; // Map one Transformers 4.57 hidden_states tuple index to the corresponding // llama.cpp graph output. Index zero (the embedding input) is intentionally not // exposed. Qwen2.5-VL has no DeepStack and its final tuple item is result_norm; -// Qwen3-VL retains the recorder/alias behavior used by the existing parity -// contract. -bool qwen_vl_hidden_state_source(QwenVLArchitecture architecture, - int decoder_layer_count, - int deepstack_layer_count, - int32_t hidden_tuple_index, - QwenVLHiddenStateSource & source, - std::string & error); +// Qwen3-VL additionally maps the DeepStack outputs exposed by llama.cpp. +bool qwen_vl_hidden_state_source(QwenVLArchitecture architecture, int decoder_layer_count, int deepstack_layer_count, + int32_t hidden_tuple_index, QwenVLHiddenStateSource & source, std::string & error); struct Qwen3VLImageView { const uint8_t * data = nullptr; - int width = 0; - int height = 0; - int channels = 0; - int stride_bytes = 0; + int width = 0; + int height = 0; + int channels = 0; + int stride_bytes = 0; }; struct Qwen3VLBridgeConfig { std::string text_path; std::string mmproj_path; std::string bundle_uuid; - int hidden_size = 0; + int hidden_size = 0; int input_embedding_size = 0; - int vocab_size = 0; + int vocab_size = 0; std::string action_token; - int32_t action_token_id = -1; - int expected_image_count = 0; - int image_min_tokens = 0; - int image_max_tokens = 0; + int32_t action_token_id = -1; + int expected_image_count = 0; + int image_min_tokens = 0; + int image_max_tokens = 0; int image_spatial_merge_size = 0; - int n_ctx = 2048; - int n_batch = 2048; - int n_threads = 0; - int verbosity = 0; - // OFT uses plain flash attention to meet its final-action parity profile. - // Other StarVLA variants retain non-flash text attention. + int n_ctx = 2048; + int n_batch = 2048; + int n_threads = 0; + int verbosity = 0; + // OFT uses flash attention; other variants require intermediate outputs + // that are only available on the non-flash path. bool flash_text_attention = false; // Round each F32 decoder residual output, plus DeepStack outputs when the // detected architecture has them, through BF16 RNE before it feeds the next @@ -97,7 +89,7 @@ struct Qwen3VLBridgeConfig { struct QwenVLGenerationConfig { size_t max_length = 0; std::vector eos_token_ids; - int top_k = 0; + int top_k = 0; float repetition_penalty = 0.0f; }; @@ -111,33 +103,28 @@ struct QwenVLGenerationResult { // generation profile: Hugging Face repetition penalty over the full sequence, // followed by top_k=1. Exposed so the generation contract can be tested // without loading a multi-gigabyte Qwen checkpoint. -bool qwen_vl_select_repetition_penalized_top1( - const float * logits, size_t vocab_size, - const std::vector & full_sequence, float repetition_penalty, - int32_t & token, std::string & error); +bool qwen_vl_select_repetition_penalized_top1(const float * logits, size_t vocab_size, + const std::vector & full_sequence, float repetition_penalty, + int32_t & token, std::string & error); class Qwen3VLBridge { public: ~Qwen3VLBridge(); - Qwen3VLBridge(const Qwen3VLBridge &) = delete; + Qwen3VLBridge(const Qwen3VLBridge &) = delete; Qwen3VLBridge & operator=(const Qwen3VLBridge &) = delete; - static std::unique_ptr load(const Qwen3VLBridgeConfig & config, - std::string & error); + static std::unique_ptr load(const Qwen3VLBridgeConfig & config, std::string & error); - bool extract_token_embeddings(const std::vector & images, - const std::string & instruction, int32_t token_id, - size_t token_count, std::vector & embeddings, + bool extract_token_embeddings(const std::vector & images, const std::string & instruction, + int32_t token_id, size_t token_count, std::vector & embeddings, std::string & error); // Full conditioning sequence. Qwen3 uses the outer recorder's raw final // decoder output (`l_out-(N-1)`); Qwen2.5 uses `result_norm`, matching its // Transformers hidden_states[-1]. Values are widened from BF16. - bool extract_full_hidden_states(const std::vector & images, - const std::string & instruction, - std::vector & hidden_states, - std::vector & attention_mask, + bool extract_full_hidden_states(const std::vector & images, const std::string & instruction, + std::vector & hidden_states, std::vector & attention_mask, std::string & error); // hidden_tuple_indices use the pinned Transformers 4.57 convention. Index @@ -146,25 +133,20 @@ class Qwen3VLBridge { // entries, including N, are raw `l_out`. For Qwen2.5, indices 1..N-1 are // raw `l_out` and index N is `result_norm`. The result is layer-major // [requested states, tokens, hidden size]. - bool extract_layer_hidden_states(const std::vector & images, - const std::string & instruction, + bool extract_layer_hidden_states(const std::vector & images, const std::string & instruction, const std::vector & hidden_tuple_indices, - std::vector & hidden_states, - std::vector & attention_mask, + std::vector & hidden_states, std::vector & attention_mask, std::string & error); // Runs a full multimodal prefill followed by incremental KV-cached text // decoding. The returned sequence includes the prompt, matching // Transformers generate(return_dict_in_generate=false). - bool generate_autoregressive( - const std::vector & images, - const std::string & instruction, - const QwenVLGenerationConfig & generation, - QwenVLGenerationResult & result, std::string & error); + bool generate_autoregressive(const std::vector & images, const std::string & instruction, + const QwenVLGenerationConfig & generation, QwenVLGenerationResult & result, + std::string & error); void reset(); const char * backend_name() const; - const char * text_attention_mode_name() const; QwenVLArchitecture architecture() const; private: @@ -175,10 +157,4 @@ class Qwen3VLBridge { std::unique_ptr impl_; }; -// Neutral aliases for new callers. The original names remain the ABI/source -// compatibility surface for the completed Qwen3 integrations. -using QwenVLImageView = Qwen3VLImageView; -using QwenVLBridgeConfig = Qwen3VLBridgeConfig; -using QwenVLBridge = Qwen3VLBridge; - } // namespace robotcpp::starvla diff --git a/src/models/starvla/qwen_bf16_round_cuda.cu b/src/models/starvla/qwen_bf16_round_cuda.cu new file mode 100644 index 0000000..2540020 --- /dev/null +++ b/src/models/starvla/qwen_bf16_round_cuda.cu @@ -0,0 +1,218 @@ +#include "models/starvla/qwen_bf16_round_cuda.h" + +#include "ggml-backend.h" +#include "ggml.h" + +#include + +#include + +namespace robotcpp::starvla { +namespace { + +__device__ uint32_t bf16_bits(float value) { + const uint32_t bits = __float_as_uint(value); + return (bits & 0x7fffffffU) > 0x7f800000U + ? (bits >> 16) | 64U + : (bits + 0x7fffU + ((bits >> 16) & 1U)) >> 16; +} + +__global__ void round_bf16(float * values, size_t count) { + const size_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (index < count) { + values[index] = __uint_as_float(bf16_bits(values[index]) << 16); + } +} + +__global__ void capture_bf16(const float * source, float * destination, + size_t count) { + const size_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (index < count) { + destination[index] = __uint_as_float(bf16_bits(source[index]) << 16); + } +} + +const char * cuda_error(cudaError_t status) { + return cudaGetErrorString(status); +} + +cudaError_t select_device(int device, int & previous_device) { + const cudaError_t status = cudaGetDevice(&previous_device); + return status == cudaSuccess ? cudaSetDevice(device) : status; +} + +QwenBF16RoundStatus tensor_device(ggml_tensor * tensor, size_t count, + int & device_id, std::string & error) { + if (tensor == nullptr || tensor->buffer == nullptr || tensor->data == nullptr || + tensor->type != GGML_TYPE_F32 || count == 0) { + error = "Qwen-VL CUDA BF16 operation received an invalid tensor"; + return QwenBF16RoundStatus::error; + } + ggml_backend_buffer_type_t buffer_type = + ggml_backend_buffer_get_type(tensor->buffer); + ggml_backend_dev_t device = + buffer_type == nullptr ? nullptr : ggml_backend_buft_get_device(buffer_type); + if (device == nullptr || + ggml_backend_dev_type(device) != GGML_BACKEND_DEVICE_TYPE_GPU) { + return QwenBF16RoundStatus::unavailable; + } + + cudaPointerAttributes attributes{}; + const cudaError_t status = cudaPointerGetAttributes(&attributes, tensor->data); + if (status == cudaErrorInvalidValue) { + cudaGetLastError(); + return QwenBF16RoundStatus::unavailable; + } + if (status != cudaSuccess) { + error = std::string("failed to inspect Qwen-VL CUDA tensor: ") + + cuda_error(status); + return QwenBF16RoundStatus::error; + } + if (attributes.type != cudaMemoryTypeDevice && + attributes.type != cudaMemoryTypeManaged) { + return QwenBF16RoundStatus::unavailable; + } + device_id = attributes.device; + return QwenBF16RoundStatus::success; +} + +} // namespace + +QwenBF16CaptureCuda::~QwenBF16CaptureCuda() { + if (data == nullptr) { + return; + } + int previous_device = 0; + if (select_device(device, previous_device) == cudaSuccess) { + cudaFree(data); + cudaSetDevice(previous_device); + } +} + +QwenBF16RoundStatus qwen_bf16_round_cuda(ggml_tensor * tensor, size_t count, + std::string & error) { + error.clear(); + int device = -1; + const QwenBF16RoundStatus available = + tensor_device(tensor, count, device, error); + if (available != QwenBF16RoundStatus::success) { + return available; + } + + int previous_device = 0; + cudaError_t status = select_device(device, previous_device); + if (status != cudaSuccess) { + error = std::string("failed to select the Qwen-VL CUDA device: ") + + cuda_error(status); + return QwenBF16RoundStatus::error; + } + + constexpr int block_size = 256; + const size_t block_count = (count + block_size - 1) / block_size; + round_bf16<<>>( + static_cast(tensor->data), count); + status = cudaGetLastError(); + if (status == cudaSuccess) { + status = cudaStreamSynchronize(cudaStreamPerThread); + } + const cudaError_t restore_status = cudaSetDevice(previous_device); + if (status != cudaSuccess) { + error = std::string("failed to round Qwen-VL residuals on CUDA: ") + + cuda_error(status); + return QwenBF16RoundStatus::error; + } + if (restore_status != cudaSuccess) { + error = std::string("failed to restore the active CUDA device: ") + + cuda_error(restore_status); + return QwenBF16RoundStatus::error; + } + return QwenBF16RoundStatus::success; +} + +QwenBF16RoundStatus qwen_bf16_capture_cuda( + ggml_tensor * tensor, size_t count, size_t offset, size_t total_count, + QwenBF16CaptureCuda & capture, std::string & error) { + error.clear(); + if (offset > total_count || count > total_count - offset || + total_count > SIZE_MAX / sizeof(float)) { + error = "Qwen-VL CUDA capture range is invalid"; + return QwenBF16RoundStatus::error; + } + int device = -1; + const QwenBF16RoundStatus available = + tensor_device(tensor, count, device, error); + if (available != QwenBF16RoundStatus::success) { + return available; + } + + int previous_device = 0; + cudaError_t status = select_device(device, previous_device); + if (status != cudaSuccess) { + error = std::string("failed to select the Qwen-VL CUDA device: ") + + cuda_error(status); + return QwenBF16RoundStatus::error; + } + if (capture.data != nullptr && + (capture.device != device || capture.capacity < total_count)) { + status = cudaFree(capture.data); + if (status == cudaSuccess) { + capture.data = nullptr; + capture.capacity = 0; + capture.device = -1; + } + } + if (status == cudaSuccess && capture.data == nullptr) { + status = cudaMalloc(&capture.data, total_count * sizeof(float)); + if (status == cudaSuccess) { + capture.capacity = total_count; + capture.device = device; + } + } + constexpr int block_size = 256; + const size_t block_count = (count + block_size - 1) / block_size; + if (status == cudaSuccess) { + capture_bf16<<>>( + static_cast(tensor->data), + static_cast(capture.data) + offset, count); + status = cudaGetLastError(); + } + if (status == cudaSuccess) { + status = cudaStreamSynchronize(cudaStreamPerThread); + } + const cudaError_t restore_status = cudaSetDevice(previous_device); + if (status != cudaSuccess || restore_status != cudaSuccess) { + error = std::string("failed to capture Qwen-VL hidden states on CUDA: ") + + cuda_error(status != cudaSuccess ? status : restore_status); + return QwenBF16RoundStatus::error; + } + return QwenBF16RoundStatus::success; +} + +bool qwen_bf16_capture_download_cuda(QwenBF16CaptureCuda & capture, + float * values, size_t count, + std::string & error) { + error.clear(); + if (capture.data == nullptr || values == nullptr || count == 0 || + count > capture.capacity || count > SIZE_MAX / sizeof(float)) { + error = "Qwen-VL CUDA capture download is invalid"; + return false; + } + int previous_device = 0; + cudaError_t status = select_device(capture.device, previous_device); + if (status != cudaSuccess) { + error = std::string("failed to select the Qwen-VL CUDA device: ") + + cuda_error(status); + return false; + } + status = cudaMemcpy(values, capture.data, count * sizeof(float), + cudaMemcpyDeviceToHost); + const cudaError_t restore_status = cudaSetDevice(previous_device); + if (status != cudaSuccess || restore_status != cudaSuccess) { + error = std::string("failed to download Qwen-VL CUDA hidden states: ") + + cuda_error(status != cudaSuccess ? status : restore_status); + return false; + } + return true; +} + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/qwen_bf16_round_cuda.h b/src/models/starvla/qwen_bf16_round_cuda.h new file mode 100644 index 0000000..985ed97 --- /dev/null +++ b/src/models/starvla/qwen_bf16_round_cuda.h @@ -0,0 +1,28 @@ +#pragma once + +#include +#include + +struct ggml_tensor; + +namespace robotcpp::starvla { + +enum class QwenBF16RoundStatus { unavailable, success, error }; + +struct QwenBF16CaptureCuda { + QwenBF16CaptureCuda() = default; + ~QwenBF16CaptureCuda(); + QwenBF16CaptureCuda(const QwenBF16CaptureCuda &) = delete; + QwenBF16CaptureCuda & operator=(const QwenBF16CaptureCuda &) = delete; + + void * data = nullptr; + size_t capacity = 0; + int device = -1; +}; + +QwenBF16RoundStatus qwen_bf16_round_cuda(ggml_tensor * tensor, size_t count, std::string & error); +QwenBF16RoundStatus qwen_bf16_capture_cuda(ggml_tensor * tensor, size_t count, size_t offset, size_t total_count, + QwenBF16CaptureCuda & capture, std::string & error); +bool qwen_bf16_capture_download_cuda(QwenBF16CaptureCuda & capture, float * values, size_t count, std::string & error); + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/starvla_engine.cpp b/src/models/starvla/starvla_engine.cpp new file mode 100644 index 0000000..db23ad5 --- /dev/null +++ b/src/models/starvla/starvla_engine.cpp @@ -0,0 +1,987 @@ +#include "models/starvla/starvla_engine.h" + +#include "ggml.h" +#include "gguf.h" +#include "models/starvla/fast_policy.h" +#include "models/starvla/groot_policy.h" +#include "models/starvla/groot_prompt.h" +#include "models/starvla/normalization.h" +#include "models/starvla/oft_image_preprocess.h" +#include "models/starvla/oft_policy.h" +#include "models/starvla/oft_prompt.h" +#include "models/starvla/pi_policy.h" +#include "models/starvla/pi_v3_policy.h" +#include "models/starvla/qwen3vl_bridge.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace robotcpp::starvla { + +const char * starvla_variant_name(StarVLAVariant variant) noexcept { + switch (variant) { + case StarVLAVariant::qwen3_oft: + return "qwen3_oft"; + case StarVLAVariant::qwen3_groot: + return "qwen3_groot"; + case StarVLAVariant::qwen3_pi_v3: + return "qwen3_pi_v3"; + case StarVLAVariant::qwen25_oft: + return "qwen25_oft"; + case StarVLAVariant::qwen25_groot: + return "qwen25_groot"; + case StarVLAVariant::qwen25_pi: + return "qwen25_pi"; + case StarVLAVariant::qwen25_fast: + return "qwen25_fast"; + } + return "unknown"; +} + +const char * starvla_variant_framework(StarVLAVariant variant) noexcept { + switch (variant) { + case StarVLAVariant::qwen3_oft: + case StarVLAVariant::qwen25_oft: + return "oft"; + case StarVLAVariant::qwen3_groot: + case StarVLAVariant::qwen25_groot: + return "groot"; + case StarVLAVariant::qwen3_pi_v3: + return "pi_v3"; + case StarVLAVariant::qwen25_pi: + return "pi"; + case StarVLAVariant::qwen25_fast: + return "fast"; + } + return "unknown"; +} + +bool starvla_variant_from_metadata(const std::string & framework, const std::string & backbone, + StarVLAVariant & variant) noexcept { + if (backbone == "qwen3_vl") { + if (framework == "oft") + variant = StarVLAVariant::qwen3_oft; + else if (framework == "groot") + variant = StarVLAVariant::qwen3_groot; + else if (framework == "pi_v3") + variant = StarVLAVariant::qwen3_pi_v3; + else + return false; + return true; + } + if (backbone == "qwen2_5_vl") { + if (framework == "oft") + variant = StarVLAVariant::qwen25_oft; + else if (framework == "groot") + variant = StarVLAVariant::qwen25_groot; + else if (framework == "pi") + variant = StarVLAVariant::qwen25_pi; + else if (framework == "fast") + variant = StarVLAVariant::qwen25_fast; + else + return false; + return true; + } + return false; +} + +namespace { + +using Clock = std::chrono::steady_clock; + +constexpr int kDefaultThreadCount = 4; + +double elapsed_ms(Clock::time_point start, Clock::time_point end) { + return std::chrono::duration(end - start).count(); +} + +bool read_policy_variant(const std::filesystem::path & path, StarVLAVariant & variant, std::string & error) { + gguf_init_params params{}; + params.no_alloc = true; + gguf_context * gguf = gguf_init_from_file(path.string().c_str(), params); + if (gguf == nullptr) { + error = "failed to read StarVLA policy GGUF metadata"; + return false; + } + const auto read_string = [&](const char * key, std::string & value) { + const int index = gguf_find_key(gguf, key); + if (index < 0 || gguf_get_kv_type(gguf, index) != GGUF_TYPE_STRING) { + error = std::string("missing StarVLA policy metadata: ") + key; + return false; + } + value = gguf_get_val_str(gguf, index); + return true; + }; + + std::string framework; + std::string backbone; + const bool valid = read_string("starvla.framework", framework) && read_string("starvla.backbone.arch", backbone); + gguf_free(gguf); + if (!valid) { + return false; + } + + if (starvla_variant_from_metadata(framework, backbone, variant)) { + return true; + } + error = "unsupported StarVLA variant: " + backbone + "/" + framework; + return false; +} + +bool is_plain_basename(const std::string & value) { + if (value.empty() || value.find('\0') != std::string::npos) { + return false; + } + const std::filesystem::path path(value); + return !path.has_root_path() && !path.has_parent_path() && path.filename() == path && value != "." && value != ".."; +} + +bool require_regular_file(const std::filesystem::path & path, const char * label, std::string & error) { + std::error_code status_error; + const bool regular = std::filesystem::is_regular_file(path, status_error); + if (!regular) { + error = std::string("StarVLA ") + label + " is not a regular file: " + path.string(); + if (status_error) { + error += " (" + status_error.message() + ")"; + } + return false; + } + return true; +} + +bool resolve_component_path(const std::string & metadata_filename, const std::string & component_path, + const char * label, std::filesystem::path & resolved, std::string & error) { + if (!is_plain_basename(metadata_filename)) { + error = std::string("StarVLA policy ") + label + " filename must be a plain basename: " + metadata_filename; + return false; + } + + if (component_path.empty() || component_path.find('\0') != std::string::npos) { + error = std::string("StarVLA ") + label + " path is required and must not contain an embedded NUL"; + return false; + } + resolved = std::filesystem::path(component_path); + if (resolved.filename().string() != metadata_filename) { + error = std::string("StarVLA ") + label + " basename must match policy metadata '" + metadata_filename + + "': " + resolved.string(); + return false; + } + return require_regular_file(resolved, label, error); +} + +bool same_file(const std::filesystem::path & lhs, const std::filesystem::path & rhs) { + std::error_code equivalent_error; + return std::filesystem::equivalent(lhs, rhs, equivalent_error) && !equivalent_error; +} + +bool validate_observation(const observation & obs, int image_count, const std::vector & image_names, + bool state_supported, const char * framework, std::string & error) { + const std::string label = std::string("StarVLA ") + framework; + if (obs.images.size() != static_cast(image_count)) { + error = label + " requires exactly " + std::to_string(image_count) + " image(s) in policy order"; + return false; + } + if (obs.task.empty()) { + error = label + " task must not be empty"; + return false; + } + if (obs.task.find('\0') != std::string::npos) { + error = label + " task contains an embedded NUL"; + return false; + } + if (!state_supported && !obs.state.empty()) { + error = label + " released checkpoint does not support state input"; + return false; + } + for (float value : obs.state) { + if (!std::isfinite(value)) { + error = label + " state must contain only finite values"; + return false; + } + } + + for (size_t i = 0; i < obs.images.size(); ++i) { + const model_image & image = obs.images[i]; + if (image.name != image_names[i]) { + error = label + " image " + std::to_string(i) + " must be named '" + image_names[i] + "'"; + return false; + } + if (image.data == nullptr || image.width <= 0 || image.height <= 0) { + error = label + " image '" + image.name + "' has invalid data or dimensions"; + return false; + } + if (image.channels != 3) { + error = label + " image '" + image.name + "' must use interleaved RGB channels"; + return false; + } + if (image.width > std::numeric_limits::max() / image.channels) { + error = label + " image '" + image.name + "' row size overflows"; + return false; + } + const int packed_stride = image.width * image.channels; + if (image.stride_bytes < 0 || (image.stride_bytes != 0 && image.stride_bytes < packed_stride)) { + error = label + " image '" + image.name + "' stride is smaller than a packed RGB row"; + return false; + } + } + return true; +} + +template +bool prepare_qwen_images(const observation & obs, const PolicyConfig & config, const char * framework, + std::vector> & processed_images, + std::vector & qwen_images, std::string & error) { + processed_images.clear(); + qwen_images.clear(); + processed_images.resize(obs.images.size()); + qwen_images.reserve(obs.images.size()); + for (size_t i = 0; i < obs.images.size(); ++i) { + const model_image & image = obs.images[i]; + int target_width = 0; + int target_height = 0; + int image_token_count = 0; + std::string preprocess_error; + if (!preprocess_qwen3vl_rgb( + image.data, image.width, image.height, image.channels, image.stride_bytes, config.image_patch_size, + config.image_spatial_merge_size, config.image_processor_min_pixels, config.image_processor_max_pixels, + processed_images[i], target_width, target_height, image_token_count, preprocess_error)) { + error = std::string("failed to preprocess StarVLA ") + framework + " image '" + image.name + + "': " + preprocess_error; + return false; + } + const uint64_t expected_bytes = static_cast(target_width) * static_cast(target_height) * 3; + if (expected_bytes != processed_images[i].size() || image_token_count < config.image_min_token_count || + image_token_count > config.image_max_token_count) { + error = std::string("StarVLA ") + framework + " image preprocessor returned an incompatible dynamic grid"; + return false; + } + Qwen3VLImageView view; + view.data = processed_images[i].data(); + view.width = target_width; + view.height = target_height; + view.channels = 3; + view.stride_bytes = target_width * 3; + qwen_images.push_back(view); + } + return true; +} + +bool prepare_pi_qwen_images(const observation & obs, const PIPolicyConfig & config, + std::vector> & pre_resized_images, + std::vector> & processed_images, + std::vector & qwen_images, std::string & error) { + pre_resized_images.clear(); + processed_images.clear(); + qwen_images.clear(); + pre_resized_images.resize(obs.images.size()); + processed_images.resize(obs.images.size()); + qwen_images.reserve(obs.images.size()); + for (size_t i = 0; i < obs.images.size(); ++i) { + const model_image & image = obs.images[i]; + std::string preprocess_error; + if (!resize_torchvision_bicubic_aa_rgb(image.data, image.width, image.height, image.stride_bytes, + config.image_framework_inference_pre_resize_width, + config.image_framework_inference_pre_resize_height, + pre_resized_images[i], preprocess_error)) { + error = "failed to pre-resize StarVLA PI image '" + image.name + "': " + preprocess_error; + return false; + } + + int target_width = 0; + int target_height = 0; + int image_token_count = 0; + if (!preprocess_qwen3vl_rgb(pre_resized_images[i].data(), config.image_framework_inference_pre_resize_width, + config.image_framework_inference_pre_resize_height, 3, + config.image_framework_inference_pre_resize_width * 3, config.image_patch_size, + config.image_spatial_merge_size, config.image_processor_min_pixels, + config.image_processor_max_pixels, processed_images[i], target_width, target_height, + image_token_count, preprocess_error)) { + error = "failed to preprocess StarVLA PI image '" + image.name + "': " + preprocess_error; + return false; + } + const uint64_t expected_bytes = static_cast(target_width) * static_cast(target_height) * 3; + if (expected_bytes != processed_images[i].size() || image_token_count < config.image_min_token_count || + image_token_count > config.image_max_token_count) { + error = "StarVLA PI image preprocessor returned an incompatible dynamic grid"; + return false; + } + Qwen3VLImageView view; + view.data = processed_images[i].data(); + view.width = target_width; + view.height = target_height; + view.channels = 3; + view.stride_bytes = target_width * 3; + qwen_images.push_back(view); + } + return true; +} + +} // namespace + +struct StarVLAEngine::Impl { + StarVLAVariant variant = StarVLAVariant::qwen3_oft; + std::filesystem::path policy_path; + std::filesystem::path text_path; + std::filesystem::path mmproj_path; + std::string normalization_profile_key; + const NormalizationConfig * normalization = nullptr; + std::mt19937_64 noise_rng; + // Destroy the policy scheduler/backends before Qwen releases llama's global backend state. + std::unique_ptr qwen; + std::unique_ptr oft_policy; + std::unique_ptr groot_policy; + std::unique_ptr pi_policy; + std::unique_ptr pi_v3_policy; + std::unique_ptr fast_policy; +}; + +StarVLAEngine::StarVLAEngine(std::unique_ptr impl) : impl_(std::move(impl)) {} + +StarVLAEngine::~StarVLAEngine() = default; + +std::unique_ptr StarVLAEngine::load(const StarVLAEngineConfig & config, std::string & error) { + error.clear(); + if (config.policy_path.empty() || config.policy_path.find('\0') != std::string::npos) { + error = "StarVLA policy path is required and must not contain an embedded NUL"; + return nullptr; + } + if (config.n_ctx <= 0 || config.n_batch <= 0 || config.n_threads < 0) { + error = "StarVLA n_ctx/n_batch must be positive and n_threads must be non-negative"; + return nullptr; + } + + std::unique_ptr impl(new Impl()); + const int effective_threads = config.n_threads > 0 ? config.n_threads : kDefaultThreadCount; + impl->policy_path = std::filesystem::path(config.policy_path); + if (!require_regular_file(impl->policy_path, "policy GGUF", error)) { + return nullptr; + } + if (!read_policy_variant(impl->policy_path, impl->variant, error)) { + return nullptr; + } + const bool is_oft = impl->variant == StarVLAVariant::qwen3_oft || impl->variant == StarVLAVariant::qwen25_oft; + const bool is_groot = impl->variant == StarVLAVariant::qwen3_groot || impl->variant == StarVLAVariant::qwen25_groot; + const bool is_pi = impl->variant == StarVLAVariant::qwen25_pi; + const bool is_pi_v3 = impl->variant == StarVLAVariant::qwen3_pi_v3; + const bool is_fast = impl->variant == StarVLAVariant::qwen25_fast; + const char * framework = starvla_variant_framework(impl->variant); + + std::string bundle_uuid; + std::string text_filename; + std::string mmproj_filename; + std::string qwen_backbone_arch; + int qwen_hidden_size = 0; + int qwen_input_embedding_size = 0; + int qwen_vocab_size = 0; + int image_count = 0; + int image_min_tokens = 0; + int image_max_tokens = 0; + int image_spatial_merge_size = 0; + const NormalizationConfig * normalization = nullptr; + if (is_oft) { + impl->oft_policy = OFTPolicy::load(impl->policy_path.string(), effective_threads, config.verbosity, error); + if (impl->oft_policy == nullptr) { + error = "failed to load StarVLA OFT policy: " + error; + return nullptr; + } + const OFTPolicyConfig & policy = impl->oft_policy->config(); + bundle_uuid = policy.bundle_uuid; + text_filename = policy.text_filename; + mmproj_filename = policy.mmproj_filename; + qwen_backbone_arch = policy.backbone_arch; + qwen_hidden_size = policy.input_dim; + qwen_input_embedding_size = policy.input_embedding_dim; + qwen_vocab_size = policy.vocab_size; + image_count = policy.image_count; + image_min_tokens = policy.image_min_token_count; + image_max_tokens = policy.image_max_token_count; + image_spatial_merge_size = policy.image_spatial_merge_size; + normalization = &policy.normalization; + } else if (is_groot) { + impl->groot_policy = GR00TPolicy::load(impl->policy_path.string(), effective_threads, config.verbosity, error); + if (impl->groot_policy == nullptr) { + error = "failed to load StarVLA GR00T policy: " + error; + return nullptr; + } + const GR00TPolicyConfig & policy = impl->groot_policy->config(); + bundle_uuid = policy.bundle_uuid; + text_filename = policy.text_filename; + mmproj_filename = policy.mmproj_filename; + qwen_backbone_arch = policy.backbone_arch; + qwen_hidden_size = policy.qwen_hidden_dim; + qwen_input_embedding_size = policy.qwen_input_embedding_dim; + qwen_vocab_size = policy.qwen_vocab_size; + image_count = policy.image_count; + image_min_tokens = policy.image_min_token_count; + image_max_tokens = policy.image_max_token_count; + image_spatial_merge_size = policy.image_spatial_merge_size; + normalization = &policy.normalization; + } else if (is_pi) { + impl->pi_policy = PIPolicy::load(impl->policy_path.string(), effective_threads, config.verbosity, error); + if (impl->pi_policy == nullptr) { + error = "failed to load StarVLA PI policy: " + error; + return nullptr; + } + const PIPolicyConfig & policy = impl->pi_policy->config(); + bundle_uuid = policy.bundle_uuid; + text_filename = policy.text_filename; + mmproj_filename = policy.mmproj_filename; + qwen_backbone_arch = policy.backbone_arch; + qwen_hidden_size = policy.qwen_hidden_dim; + qwen_input_embedding_size = policy.qwen_input_embedding_dim; + qwen_vocab_size = policy.qwen_vocab_size; + image_count = policy.image_count; + image_min_tokens = policy.image_min_token_count; + image_max_tokens = policy.image_max_token_count; + image_spatial_merge_size = policy.image_spatial_merge_size; + normalization = &policy.normalization; + } else if (is_pi_v3) { + impl->pi_v3_policy = PIV3Policy::load(impl->policy_path.string(), effective_threads, config.verbosity, error); + if (impl->pi_v3_policy == nullptr) { + error = "failed to load StarVLA PI_v3 policy: " + error; + return nullptr; + } + const PIV3PolicyConfig & policy = impl->pi_v3_policy->config(); + bundle_uuid = policy.bundle_uuid; + text_filename = policy.text_filename; + mmproj_filename = policy.mmproj_filename; + qwen_backbone_arch = policy.backbone_arch; + qwen_hidden_size = policy.qwen_hidden_dim; + qwen_input_embedding_size = policy.qwen_input_embedding_dim; + qwen_vocab_size = policy.qwen_vocab_size; + image_count = policy.image_count; + image_min_tokens = policy.image_min_token_count; + image_max_tokens = policy.image_max_token_count; + image_spatial_merge_size = policy.image_spatial_merge_size; + normalization = &policy.normalization; + } else { + impl->fast_policy = FastPolicy::load(impl->policy_path.string(), config.verbosity, error); + if (impl->fast_policy == nullptr) { + error = "failed to load StarVLA FAST policy: " + error; + return nullptr; + } + const FastPolicyConfig & policy = impl->fast_policy->config(); + bundle_uuid = policy.bundle_uuid; + text_filename = policy.text_filename; + mmproj_filename = policy.mmproj_filename; + qwen_backbone_arch = policy.backbone_arch; + qwen_hidden_size = policy.qwen_hidden_dim; + qwen_input_embedding_size = policy.qwen_input_embedding_dim; + qwen_vocab_size = policy.qwen_vocab_size; + image_count = policy.image_count; + image_min_tokens = policy.image_min_token_count; + image_max_tokens = policy.image_max_token_count; + image_spatial_merge_size = policy.image_spatial_merge_size; + normalization = &policy.normalization; + if (config.n_ctx < static_cast(policy.generation_max_length)) { + error = "StarVLA FAST --n-ctx must be at least max_length=" + std::to_string(policy.generation_max_length); + return nullptr; + } + } + if (!resolve_component_path(text_filename, config.text_path, "text GGUF", impl->text_path, error) || + !resolve_component_path(mmproj_filename, config.mmproj_path, "mmproj GGUF", impl->mmproj_path, error)) { + return nullptr; + } + if (same_file(impl->policy_path, impl->text_path) || same_file(impl->policy_path, impl->mmproj_path) || + same_file(impl->text_path, impl->mmproj_path)) { + error = "StarVLA policy, text, and mmproj GGUF paths must identify three distinct files"; + return nullptr; + } + + impl->normalization = normalization; + std::string profile_error; + const NormalizationProfile * profile = resolve_normalization_profile(*normalization, "", profile_error); + if (profile == nullptr) { + error = std::string("failed to select StarVLA ") + framework + " normalization profile: " + profile_error; + return nullptr; + } + impl->normalization_profile_key = profile->key; + + Qwen3VLBridgeConfig qwen_config; + qwen_config.text_path = impl->text_path.string(); + qwen_config.mmproj_path = impl->mmproj_path.string(); + qwen_config.bundle_uuid = bundle_uuid; + qwen_config.hidden_size = qwen_hidden_size; + qwen_config.input_embedding_size = qwen_input_embedding_size; + qwen_config.vocab_size = qwen_vocab_size; + if (is_oft) { + qwen_config.action_token = impl->oft_policy->config().prompt.action_token; + qwen_config.action_token_id = impl->oft_policy->config().action_token_id; + } else { + qwen_config.action_token.clear(); + qwen_config.action_token_id = -1; + } + qwen_config.expected_image_count = image_count; + qwen_config.image_min_tokens = image_min_tokens; + qwen_config.image_max_tokens = image_max_tokens; + qwen_config.image_spatial_merge_size = image_spatial_merge_size; + qwen_config.n_ctx = config.n_ctx; + qwen_config.n_batch = config.n_batch; + qwen_config.n_threads = effective_threads; + qwen_config.verbosity = config.verbosity; + qwen_config.flash_text_attention = is_oft || is_pi; + qwen_config.bf16_residual_layer_boundaries = is_groot; + qwen_config.disable_text_backend_native_graphs = true; + qwen_config.disable_vision_backend_native_graphs = true; + if (qwen_config.input_embedding_size <= 0) { + error = "StarVLA Qwen-VL input embedding size is invalid"; + return nullptr; + } + impl->qwen = Qwen3VLBridge::load(qwen_config, error); + if (impl->qwen == nullptr) { + error = "failed to load StarVLA Qwen-VL components: " + error; + return nullptr; + } + const QwenVLArchitecture expected_architecture = + qwen_backbone_arch == "qwen2_5_vl" ? QwenVLArchitecture::qwen2_5_vl : QwenVLArchitecture::qwen3_vl; + if (impl->qwen->architecture() != expected_architecture) { + error = "StarVLA policy backbone metadata does not match the Qwen-VL components"; + return nullptr; + } + + if (!is_oft && !is_fast) { + if (config.noise_seed >= 0) { + impl->noise_rng.seed(static_cast(config.noise_seed)); + } else { + std::random_device device; + std::seed_seq seed{device(), device(), device(), device(), + static_cast(Clock::now().time_since_epoch().count())}; + impl->noise_rng.seed(seed); + } + } + + if (config.verbosity >= 1) { + const char * variant_name = starvla_variant_name(impl->variant); + const char * policy_backend = is_oft ? impl->oft_policy->backend_name() + : (is_groot ? impl->groot_policy->backend_name() + : (is_pi ? impl->pi_policy->backend_name() + : (is_pi_v3 ? impl->pi_v3_policy->backend_name() + : impl->fast_policy->backend_name()))); + std::fprintf(stderr, + "%s: variant=%s policy=%s text=%s mmproj=%s profile=%s " + "qwen_backend=%s policy_backend=%s\n", + __func__, variant_name, impl->policy_path.string().c_str(), impl->text_path.string().c_str(), + impl->mmproj_path.string().c_str(), impl->normalization_profile_key.c_str(), + impl->qwen->backend_name(), policy_backend); + } + return std::unique_ptr(new StarVLAEngine(std::move(impl))); +} + +bool StarVLAEngine::predict(const observation & obs, StarVLAEngineResult & result, std::string & error) { + result = StarVLAEngineResult{}; + error.clear(); + const Clock::time_point total_start = Clock::now(); + const auto fail = [&]() { + result.actions.clear(); + result.timings.total_ms = elapsed_ms(total_start, Clock::now()); + return false; + }; + + if (impl_ == nullptr || impl_->qwen == nullptr || + (!impl_->oft_policy && !impl_->groot_policy && !impl_->pi_policy && !impl_->pi_v3_policy && + !impl_->fast_policy)) { + error = "StarVLA engine is not initialized"; + return fail(); + } + + if (impl_->normalization == nullptr) { + error = "StarVLA normalization metadata is not initialized"; + return fail(); + } + std::string profile_error; + const NormalizationProfile * profile = + resolve_normalization_profile(*impl_->normalization, impl_->normalization_profile_key, profile_error); + if (profile == nullptr) { + error = "failed to select StarVLA normalization profile: " + profile_error; + return fail(); + } + std::vector normalized_actions; + std::string instruction; + + const auto make_noise = [&](size_t count, std::vector & noise) { + if (!obs.initial_noise.empty()) { + if (obs.initial_noise.size() != count || !std::all_of(obs.initial_noise.begin(), obs.initial_noise.end(), + [](float value) { return std::isfinite(value); })) { + error = "initial noise has an incompatible shape or non-finite value"; + return false; + } + noise = obs.initial_noise; + return true; + } + noise.resize(count); + std::normal_distribution normal(0.0f, 1.0f); + for (float & value : noise) { + value = ggml_bf16_to_fp32(ggml_fp32_to_bf16(normal(impl_->noise_rng))); + } + return true; + }; + + if (impl_->variant == StarVLAVariant::qwen25_fast) { + if (!obs.initial_noise.empty()) { + error = "StarVLA FAST does not use diffusion noise"; + return fail(); + } + const FastPolicyConfig & config = impl_->fast_policy->config(); + if (!validate_observation(obs, config.image_count, config.image_names, false, "FAST", error)) { + return fail(); + } + + Clock::time_point stage_start = Clock::now(); + std::vector> processed_images; + std::vector qwen_images; + if (!prepare_qwen_images(obs, config, "FAST", processed_images, qwen_images, error)) { + return fail(); + } + result.timings.image_preprocess_ms = elapsed_ms(stage_start, Clock::now()); + + stage_start = Clock::now(); + if (!build_fast_instruction(config.cot_template, obs.task, instruction, error)) { + error = "failed to build the StarVLA FAST prompt: " + error; + return fail(); + } + result.timings.prompt_ms = elapsed_ms(stage_start, Clock::now()); + + QwenVLGenerationConfig generation; + generation.max_length = config.generation_max_length; + generation.eos_token_ids = config.generation_eos_token_ids; + generation.top_k = config.generation_top_k; + generation.repetition_penalty = config.generation_repetition_penalty; + QwenVLGenerationResult generated; + stage_start = Clock::now(); + if (!impl_->qwen->generate_autoregressive(qwen_images, instruction, generation, generated, error)) { + error = "StarVLA FAST Qwen2.5-VL generation failed: " + error; + return fail(); + } + result.timings.qwen3vl_ms = elapsed_ms(stage_start, Clock::now()); + if (generated.prompt_token_count == 0 || generated.full_sequence.size() < generated.prompt_token_count || + generated.full_sequence.size() > config.generation_max_length) { + error = "StarVLA FAST Qwen2.5-VL returned an incompatible generated sequence"; + return fail(); + } + stage_start = Clock::now(); + if (!impl_->fast_policy->decode_generated(generated.full_sequence, normalized_actions, error)) { + error = "StarVLA FAST codec decode failed: " + error; + return fail(); + } + result.timings.policy_ms = elapsed_ms(stage_start, Clock::now()); + + stage_start = Clock::now(); + if (!impl_->fast_policy->unnormalize(normalized_actions, profile->key, result.actions, error)) { + error = "StarVLA FAST action unnormalization failed: " + error; + return fail(); + } + result.timings.unnormalize_ms = elapsed_ms(stage_start, Clock::now()); + const size_t expected_actions = static_cast(config.horizon) * config.action_dim; + if (result.actions.size() != expected_actions || normalized_actions.size() != expected_actions || + !std::all_of(result.actions.begin(), result.actions.end(), + [](float action) { return std::isfinite(action); })) { + error = "StarVLA FAST returned an incompatible or non-finite action tensor"; + return fail(); + } + result.chunk_size = config.horizon; + result.action_dim = config.action_dim; + result.timings.total_ms = elapsed_ms(total_start, Clock::now()); + return true; + } + + if (impl_->variant == StarVLAVariant::qwen25_pi) { + const PIPolicyConfig & config = impl_->pi_policy->config(); + if (!validate_observation(obs, config.image_count, config.image_names, true, "PI", error)) { + return fail(); + } + if (!obs.state.empty() && obs.state.size() != static_cast(config.state_dim)) { + error = + "StarVLA PI accepts either no state or exactly " + std::to_string(config.state_dim) + " state values"; + return fail(); + } + + Clock::time_point stage_start = Clock::now(); + std::vector> pre_resized_images; + std::vector> processed_images; + std::vector qwen_images; + if (!prepare_pi_qwen_images(obs, config, pre_resized_images, processed_images, qwen_images, error)) { + return fail(); + } + result.timings.image_preprocess_ms = elapsed_ms(stage_start, Clock::now()); + + stage_start = Clock::now(); + if (!build_pi_v3_instruction(config.cot_template, obs.task, instruction, error)) { + error = "failed to build the StarVLA PI prompt: " + error; + return fail(); + } + result.timings.prompt_ms = elapsed_ms(stage_start, Clock::now()); + + std::vector hidden_states; + std::vector attention_mask; + stage_start = Clock::now(); + if (!impl_->qwen->extract_layer_hidden_states(qwen_images, instruction, config.qwen_hidden_tuple_indices, + hidden_states, attention_mask, error)) { + error = "StarVLA PI Qwen2.5-VL inference failed: " + error; + return fail(); + } + result.timings.qwen3vl_ms = elapsed_ms(stage_start, Clock::now()); + const size_t expected_hidden = static_cast(config.block_count) * attention_mask.size() * + static_cast(config.qwen_hidden_dim); + if (attention_mask.empty() || hidden_states.size() != expected_hidden) { + error = "StarVLA PI Qwen2.5-VL returned an incompatible layerwise conditioning shape"; + return fail(); + } + + std::vector noise; + if (!make_noise(static_cast(config.horizon) * config.action_dim, noise)) { + return fail(); + } + + stage_start = Clock::now(); + if (!impl_->pi_policy->evaluate(hidden_states.data(), hidden_states.size(), + obs.state.empty() ? nullptr : obs.state.data(), obs.state.size(), noise.data(), + noise.size(), normalized_actions, error)) { + error = "StarVLA PI policy inference failed: " + error; + return fail(); + } + result.timings.policy_ms = elapsed_ms(stage_start, Clock::now()); + + stage_start = Clock::now(); + if (!impl_->pi_policy->unnormalize(normalized_actions, profile->key, result.actions, error)) { + error = "StarVLA PI action unnormalization failed: " + error; + return fail(); + } + result.timings.unnormalize_ms = elapsed_ms(stage_start, Clock::now()); + const size_t expected_actions = static_cast(config.horizon) * config.action_dim; + if (result.actions.size() != expected_actions || + !std::all_of(result.actions.begin(), result.actions.end(), + [](float action) { return std::isfinite(action); })) { + error = "StarVLA PI returned an incompatible or non-finite action tensor"; + return fail(); + } + result.chunk_size = config.horizon; + result.action_dim = config.action_dim; + result.timings.total_ms = elapsed_ms(total_start, Clock::now()); + return true; + } + + if (impl_->variant == StarVLAVariant::qwen3_pi_v3) { + const PIV3PolicyConfig & config = impl_->pi_v3_policy->config(); + if (!validate_observation(obs, config.image_count, config.image_names, false, "PI_v3", error)) { + return fail(); + } + + Clock::time_point stage_start = Clock::now(); + std::vector> processed_images; + std::vector qwen_images; + if (!prepare_qwen_images(obs, config, "PI_v3", processed_images, qwen_images, error)) { + return fail(); + } + result.timings.image_preprocess_ms = elapsed_ms(stage_start, Clock::now()); + + stage_start = Clock::now(); + if (!build_pi_v3_instruction(config.cot_template, obs.task, instruction, error)) { + error = "failed to build the StarVLA PI_v3 prompt: " + error; + return fail(); + } + result.timings.prompt_ms = elapsed_ms(stage_start, Clock::now()); + + std::vector hidden_states; + std::vector attention_mask; + stage_start = Clock::now(); + if (!impl_->qwen->extract_layer_hidden_states(qwen_images, instruction, config.qwen_hidden_tuple_indices, + hidden_states, attention_mask, error)) { + error = "StarVLA PI_v3 Qwen3-VL inference failed: " + error; + return fail(); + } + result.timings.qwen3vl_ms = elapsed_ms(stage_start, Clock::now()); + const size_t expected_hidden = + static_cast(config.qwen_layer_count) * attention_mask.size() * config.qwen_hidden_dim; + if (attention_mask.empty() || hidden_states.size() != expected_hidden) { + error = "StarVLA PI_v3 Qwen3-VL returned an incompatible layerwise conditioning shape"; + return fail(); + } + + std::vector noise; + if (!make_noise(static_cast(config.horizon) * config.action_dim, noise)) { + return fail(); + } + + stage_start = Clock::now(); + if (!impl_->pi_v3_policy->evaluate(hidden_states.data(), hidden_states.size(), attention_mask.data(), + attention_mask.size(), noise.data(), noise.size(), normalized_actions, + error)) { + error = "StarVLA PI_v3 policy inference failed: " + error; + return fail(); + } + result.timings.policy_ms = elapsed_ms(stage_start, Clock::now()); + + stage_start = Clock::now(); + if (!impl_->pi_v3_policy->unnormalize(normalized_actions, profile->key, result.actions, error)) { + error = "StarVLA PI_v3 action unnormalization failed: " + error; + return fail(); + } + result.timings.unnormalize_ms = elapsed_ms(stage_start, Clock::now()); + const size_t expected_actions = static_cast(config.horizon) * config.action_dim; + if (result.actions.size() != expected_actions || + !std::all_of(result.actions.begin(), result.actions.end(), + [](float action) { return std::isfinite(action); })) { + error = "StarVLA PI_v3 returned an incompatible or non-finite action tensor"; + return fail(); + } + result.chunk_size = config.horizon; + result.action_dim = config.action_dim; + result.timings.total_ms = elapsed_ms(total_start, Clock::now()); + return true; + } + + if (impl_->variant == StarVLAVariant::qwen3_groot || impl_->variant == StarVLAVariant::qwen25_groot) { + const GR00TPolicyConfig & config = impl_->groot_policy->config(); + if (!validate_observation(obs, config.image_count, config.image_names, false, "GR00T", error)) { + return fail(); + } + + Clock::time_point stage_start = Clock::now(); + std::vector> processed_images; + std::vector qwen_images; + if (!prepare_qwen_images(obs, config, "GR00T", processed_images, qwen_images, error)) { + return fail(); + } + result.timings.image_preprocess_ms = elapsed_ms(stage_start, Clock::now()); + + stage_start = Clock::now(); + if (!build_groot_instruction(config.cot_template, obs.task, instruction, error)) { + error = "failed to build the StarVLA GR00T prompt: " + error; + return fail(); + } + result.timings.prompt_ms = elapsed_ms(stage_start, Clock::now()); + + std::vector hidden_states; + std::vector attention_mask; + stage_start = Clock::now(); + if (!impl_->qwen->extract_full_hidden_states(qwen_images, instruction, hidden_states, attention_mask, error)) { + error = "StarVLA GR00T Qwen3-VL inference failed: " + error; + return fail(); + } + result.timings.qwen3vl_ms = elapsed_ms(stage_start, Clock::now()); + if (hidden_states.empty() || hidden_states.size() % static_cast(config.qwen_hidden_dim) != 0 || + hidden_states.size() / static_cast(config.qwen_hidden_dim) != attention_mask.size()) { + error = "StarVLA GR00T Qwen3-VL returned an incompatible conditioning shape"; + return fail(); + } + + std::vector noise; + if (!make_noise(static_cast(config.horizon) * config.action_dim, noise)) { + return fail(); + } + + stage_start = Clock::now(); + if (!impl_->groot_policy->evaluate(hidden_states.data(), hidden_states.size(), attention_mask.data(), + attention_mask.size(), noise.data(), noise.size(), normalized_actions, + error)) { + error = "StarVLA GR00T policy inference failed: " + error; + return fail(); + } + result.timings.policy_ms = elapsed_ms(stage_start, Clock::now()); + + stage_start = Clock::now(); + if (!impl_->groot_policy->unnormalize(normalized_actions, profile->key, result.actions, error)) { + error = "StarVLA GR00T action unnormalization failed: " + error; + return fail(); + } + result.timings.unnormalize_ms = elapsed_ms(stage_start, Clock::now()); + const size_t expected_actions = static_cast(config.horizon) * config.action_dim; + if (result.actions.size() != expected_actions || + !std::all_of(result.actions.begin(), result.actions.end(), + [](float action) { return std::isfinite(action); })) { + error = "StarVLA GR00T returned an incompatible or non-finite action tensor"; + return fail(); + } + result.chunk_size = config.horizon; + result.action_dim = config.action_dim; + result.timings.total_ms = elapsed_ms(total_start, Clock::now()); + return true; + } + + if (!obs.initial_noise.empty()) { + error = "StarVLA OFT does not use diffusion noise"; + return fail(); + } + const OFTPolicyConfig & config = impl_->oft_policy->config(); + if (!validate_observation(obs, config.image_count, config.image_names, true, "OFT", error)) { + return fail(); + } + + Clock::time_point stage_start = Clock::now(); + std::vector> processed_images; + std::vector qwen_images; + if (!prepare_qwen_images(obs, config, "OFT", processed_images, qwen_images, error)) { + return fail(); + } + result.timings.image_preprocess_ms = elapsed_ms(stage_start, Clock::now()); + + stage_start = Clock::now(); + if (!build_oft_instruction(config.prompt, obs.task, obs.state, instruction, error)) { + error = "failed to build the StarVLA OFT prompt: " + error; + return fail(); + } + result.timings.prompt_ms = elapsed_ms(stage_start, Clock::now()); + + std::vector action_queries; + stage_start = Clock::now(); + if (!impl_->qwen->extract_token_embeddings(qwen_images, instruction, config.action_token_id, + static_cast(config.horizon), action_queries, error)) { + error = "StarVLA OFT Qwen3-VL inference failed: " + error; + return fail(); + } + result.timings.qwen3vl_ms = elapsed_ms(stage_start, Clock::now()); + const size_t expected_queries = static_cast(config.horizon) * config.input_dim; + if (action_queries.size() != expected_queries) { + error = "StarVLA OFT Qwen3-VL returned an incompatible action-query shape"; + return fail(); + } + + stage_start = Clock::now(); + if (!impl_->oft_policy->evaluate(action_queries.data(), action_queries.size(), normalized_actions, error)) { + error = "StarVLA OFT policy inference failed: " + error; + return fail(); + } + result.timings.policy_ms = elapsed_ms(stage_start, Clock::now()); + + stage_start = Clock::now(); + if (!impl_->oft_policy->unnormalize(normalized_actions, profile->key, result.actions, error)) { + error = "StarVLA OFT action unnormalization failed: " + error; + return fail(); + } + result.timings.unnormalize_ms = elapsed_ms(stage_start, Clock::now()); + const size_t expected_actions = static_cast(config.horizon) * config.action_dim; + if (result.actions.size() != expected_actions) { + error = "StarVLA OFT returned an incompatible action tensor shape"; + return fail(); + } + for (float action : result.actions) { + if (!std::isfinite(action)) { + error = "StarVLA OFT returned a non-finite unnormalized action"; + return fail(); + } + } + + result.chunk_size = config.horizon; + result.action_dim = config.action_dim; + result.timings.total_ms = elapsed_ms(total_start, Clock::now()); + return true; +} + +void StarVLAEngine::reset() { + if (impl_ != nullptr && impl_->qwen != nullptr) { + impl_->qwen->reset(); + } +} + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/starvla_engine.h b/src/models/starvla/starvla_engine.h new file mode 100644 index 0000000..92c7319 --- /dev/null +++ b/src/models/starvla/starvla_engine.h @@ -0,0 +1,74 @@ +#pragma once + +#include "models/model.h" + +#include +#include +#include +#include + +namespace robotcpp::starvla { + +enum class StarVLAVariant { + qwen3_oft, + qwen3_groot, + qwen3_pi_v3, + qwen25_oft, + qwen25_groot, + qwen25_pi, + qwen25_fast, +}; + +const char * starvla_variant_name(StarVLAVariant variant) noexcept; +const char * starvla_variant_framework(StarVLAVariant variant) noexcept; +bool starvla_variant_from_metadata(const std::string & framework, const std::string & backbone, + StarVLAVariant & variant) noexcept; + +struct StarVLAEngineConfig { + std::string policy_path; + std::string text_path; + std::string mmproj_path; + int n_threads = 0; + int n_ctx = 2048; + int n_batch = 512; + int64_t noise_seed = -1; + int verbosity = 0; +}; + +struct StarVLAStageTimings { + double image_preprocess_ms = 0.0; + double prompt_ms = 0.0; + double qwen3vl_ms = 0.0; + double policy_ms = 0.0; + double unnormalize_ms = 0.0; + double total_ms = 0.0; +}; + +struct StarVLAEngineResult { + std::vector actions; + int chunk_size = 0; + int action_dim = 0; + StarVLAStageTimings timings; +}; + +class StarVLAEngine { + public: + ~StarVLAEngine(); + + StarVLAEngine(const StarVLAEngine &) = delete; + StarVLAEngine & operator=(const StarVLAEngine &) = delete; + + static std::unique_ptr load(const StarVLAEngineConfig & config, std::string & error); + + bool predict(const observation & obs, StarVLAEngineResult & result, std::string & error); + void reset(); + + private: + struct Impl; + + explicit StarVLAEngine(std::unique_ptr impl); + + std::unique_ptr impl_; +}; + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/starvla_model.cpp b/src/models/starvla/starvla_model.cpp new file mode 100644 index 0000000..e320676 --- /dev/null +++ b/src/models/starvla/starvla_model.cpp @@ -0,0 +1,93 @@ +#include "models/starvla/starvla_model.h" + +#include "models/starvla/starvla_engine.h" + +#include +#include +#include + +namespace robotcpp { +namespace { + +void add_metric(model_result & out, const char * name, double value) { + model_metric metric; + metric.name = name; + metric.value = value; + out.metrics.push_back(std::move(metric)); +} + +} // namespace + +StarVLAModel::StarVLAModel(std::unique_ptr engine) : engine_(std::move(engine)) {} + +StarVLAModel::~StarVLAModel() = default; + +const char * StarVLAModel::type() const { + return "starvla"; +} + +bool StarVLAModel::predict(const observation & obs, model_result & out, std::string & error) { + out = model_result{}; + error.clear(); + if (engine_ == nullptr) { + error = "StarVLA model is not initialized"; + return false; + } + + starvla::StarVLAEngineResult result; + if (!engine_->predict(obs, result, error)) { + return false; + } + out.actions = std::move(result.actions); + out.chunk_size = result.chunk_size; + out.action_dim = result.action_dim; + add_metric(out, "image_preprocess_ms", result.timings.image_preprocess_ms); + add_metric(out, "prompt_ms", result.timings.prompt_ms); + add_metric(out, "qwen3vl_ms", result.timings.qwen3vl_ms); + add_metric(out, "policy_ms", result.timings.policy_ms); + add_metric(out, "unnormalize_ms", result.timings.unnormalize_ms); + add_metric(out, "model_total_ms", result.timings.total_ms); + return true; +} + +void StarVLAModel::reset() { + if (engine_ != nullptr) { + engine_->reset(); + } +} + +bool make_starvla_model(const model_args & args, std::unique_ptr & out, std::string & error) { + out.reset(); + error.clear(); + if (!is_starvla_model_type(args.type)) { + error = std::string("model type '") + model_type_name(args.type) + "' is not a StarVLA model type"; + return false; + } + if (args.noise_mode != 0) { + error = "StarVLA does not support SmolVLA --noise-mode debug-sin; use Gaussian noise and --noise-seed"; + return false; + } + starvla::StarVLAEngineConfig config; + config.policy_path = args.policy_path; + config.text_path = args.llm_path; + config.mmproj_path = args.mmproj_path; + config.n_threads = args.threads; + config.n_ctx = args.n_ctx; + config.n_batch = args.n_batch; + config.noise_seed = args.noise_seed; + config.verbosity = args.verbosity; + std::unique_ptr engine = starvla::StarVLAEngine::load(config, error); + if (engine == nullptr) { + return false; + } + + std::unique_ptr model(new (std::nothrow) StarVLAModel(std::move(engine))); + if (model == nullptr) { + error = "failed to allocate StarVLA model"; + return false; + } + out = std::move(model); + return true; +} + +} // namespace robotcpp diff --git a/src/models/starvla/starvla_model.h b/src/models/starvla/starvla_model.h new file mode 100644 index 0000000..9c41099 --- /dev/null +++ b/src/models/starvla/starvla_model.h @@ -0,0 +1,35 @@ +#pragma once + +#include "models/model.h" + +#include +#include + +namespace robotcpp::starvla { +class StarVLAEngine; +} + +namespace robotcpp { + +class StarVLAModel final : public Model { + public: + ~StarVLAModel() override; + + StarVLAModel(const StarVLAModel &) = delete; + StarVLAModel & operator=(const StarVLAModel &) = delete; + + const char * type() const override; + bool predict(const observation & obs, model_result & out, std::string & error) override; + void reset() override; + + private: + explicit StarVLAModel(std::unique_ptr engine); + + friend bool make_starvla_model(const model_args & args, std::unique_ptr & out, std::string & error); + + std::unique_ptr engine_; +}; + +bool make_starvla_model(const model_args & args, std::unique_ptr & out, std::string & error); + +} // namespace robotcpp diff --git a/tests/starvla/fast_codec_test.cpp b/tests/starvla/fast_codec_test.cpp deleted file mode 100644 index 701d119..0000000 --- a/tests/starvla/fast_codec_test.cpp +++ /dev/null @@ -1,194 +0,0 @@ -#include "models/starvla/fast_codec.h" - -#include -#include -#include -#include -#include -#include -#include -#include - -namespace { - -using robotcpp::starvla::FastCodec; -using robotcpp::starvla::FastCodecConfig; -using robotcpp::starvla::FastDecodeResult; - -void require(bool condition, const std::string & message) { - if (!condition) { - std::cerr << "FAIL: " << message << '\n'; - std::exit(1); - } -} - -uint32_t byte_level_codepoint(uint8_t target) { - auto is_direct = [](int value) { - return (value >= 0x21 && value <= 0x7e) || (value >= 0xa1 && value <= 0xac) || - (value >= 0xae && value <= 0xff); - }; - if (is_direct(target)) { - return target; - } - uint32_t extra = 0; - for (int value = 0; value < target; ++value) { - if (!is_direct(value)) { - ++extra; - } - } - return 256U + extra; -} - -std::string utf8(uint32_t codepoint) { - std::string output; - if (codepoint <= 0x7fU) { - output.push_back(static_cast(codepoint)); - } else if (codepoint <= 0x7ffU) { - output.push_back(static_cast(0xc0U | (codepoint >> 6U))); - output.push_back(static_cast(0x80U | (codepoint & 0x3fU))); - } else { - output.push_back(static_cast(0xe0U | (codepoint >> 12U))); - output.push_back(static_cast(0x80U | ((codepoint >> 6U) & 0x3fU))); - output.push_back(static_cast(0x80U | (codepoint & 0x3fU))); - } - return output; -} - -std::unique_ptr make_synthetic_codec(size_t time_horizon = 2, - size_t action_dim = 2) { - const std::vector raw_bytes = {10, 20, 30, 40, 0xe2, 0x82, 0x28}; - std::vector vocab; - for (uint8_t byte : raw_bytes) { - vocab.push_back(utf8(byte_level_codepoint(byte))); - } - FastCodecConfig config; - config.scale = 1.0; - config.min_token = 0; - config.vocab_size = vocab.size(); - config.time_horizon = time_horizon; - config.action_dim = action_dim; - std::string error; - auto codec = FastCodec::create(config, vocab, {100, 42, 999, 7, 501, 502, 503}, error); - require(codec != nullptr, "synthetic FAST codec must construct: " + error); - return codec; -} - -void run_unit_tests() { - auto codec = make_synthetic_codec(); - std::string error; - - std::vector vlm_ids; - require(codec->map_fast_to_vlm({3, 0, 2, 1}, vlm_ids, error), - "FAST-to-VLM mapping must succeed"); - require(vlm_ids == std::vector({7, 100, 999, 42}), - "FAST-to-VLM mapping must use the explicit non-contiguous table"); - - std::vector fast_ids; - require(codec->map_vlm_to_fast(vlm_ids, fast_ids, error), - "VLM-to-FAST mapping must succeed"); - require(fast_ids == std::vector({3, 0, 2, 1}), - "VLM-to-FAST mapping must invert the explicit table"); - require(!codec->map_vlm_to_fast({101}, fast_ids, error), - "an unmapped Qwen token must be rejected"); - - require(codec->extract_fast_tokens({-1, 42, 1234, 100, 7, 42}, fast_ids, error), - "full Qwen sequence action extraction must succeed"); - require(fast_ids == std::vector({1, 0, 3, 1}), - "action extraction must filter with the inverse map and preserve order"); - - std::vector codepoints; - require(codec->byte_level_decode({0, 1, 2, 3}, codepoints, error), - "ByteLevel decode must succeed"); - require(codepoints == std::vector({10, 20, 30, 40}), - "ByteLevel decode must invert the GPT-2 byte alphabet"); - require(codec->byte_level_decode({4, 5, 6}, codepoints, error), - "lossy UTF-8 ByteLevel decode must succeed"); - require(codepoints == std::vector({0xfffdU, 0x28U}), - "ByteLevel decode must match Rust UTF-8 replacement semantics"); - - FastDecodeResult decoded; - require(codec->decode_fast_tokens({{0, 1, 2, 3}}, decoded, error), - "synthetic inverse DCT must succeed"); - require(decoded.actions.size() == 4, - "valid synthetic tokens must produce one 2x2 action chunk"); - const double root_half = std::sqrt(0.5); - const std::vector expected = { - root_half * (10.0 + 30.0), root_half * (20.0 + 40.0), - root_half * (10.0 - 30.0), root_half * (20.0 - 40.0), - }; - for (size_t index = 0; index < expected.size(); ++index) { - require(std::abs(decoded.actions[index] - expected[index]) < 1e-12, - "orthonormal inverse DCT must match the analytical result"); - } - - require(!codec->decode_fast_tokens({{}, {0, 1, 2}, {9999}}, decoded, error), - "malformed FAST coefficients must fail"); - require(decoded.actions.empty(), "failed FAST decode must not return zero actions"); - require(!codec->decode_fast_tokens({}, decoded, error), "an empty batch must fail explicitly"); - - require(!codec->decode_fast_tokens( - {{0, 1, 2, 3}, {9999}, {0, 1, 2, 3}}, decoded, error), - "a malformed FAST batch member must fail the batch"); - - FastDecodeResult decoded_generated; - require(codec->decode_generated_tokens( - {{-1, 100, 123456, 42, 999, 555555, 7}}, decoded_generated, error), - "complete generated_ids must filter then decode"); - bool generated_matches_expected = decoded_generated.actions.size() == expected.size(); - for (size_t index = 0; generated_matches_expected && index < expected.size(); ++index) { - generated_matches_expected = - std::abs(decoded_generated.actions[index] - expected[index]) < 1e-12; - } - require(generated_matches_expected, - "complete generated_ids must match the pure FAST-token action decode"); - require(!codec->decode_generated_tokens({{1, 2, 3}}, decoded_generated, error), - "a generated sequence without action tokens must fail"); - require(!codec->decode_vlm_action_tokens({{123456}}, decoded_generated, error), - "strict low-level VLM action-token decode must reject ordinary Qwen tokens"); - - std::vector maximum_generated_sequence(2048, 123456); - require(!codec->decode_generated_tokens( - {maximum_generated_sequence}, decoded_generated, error), - "max_length text without action tokens must fail decode"); - std::vector oversized_sequence(2049, 100); - require(!codec->extract_fast_tokens(oversized_sequence, fast_ids, error), - "generated token sequences beyond official max_length must fail before allocation"); - - FastCodecConfig invalid_config = codec->config(); - invalid_config.scale = 0.0; - auto invalid = FastCodec::create(invalid_config, - {"a", "b", "c", "d", "e", "f", "g"}, - {0, 1, 2, 3, 4, 5, 6}, error); - require(invalid == nullptr, "zero FAST scale must be rejected"); - - invalid_config = codec->config(); - invalid = FastCodec::create(invalid_config, - {"a", "b", "c", "d", "e", "f", "g"}, - {0, 1, 2, 3, 4, 5, 5}, error); - require(invalid == nullptr, "duplicate VLM action-token IDs must be rejected"); - - invalid_config = codec->config(); - invalid_config.time_horizon = 1025; - invalid = FastCodec::create(invalid_config, - {"a", "b", "c", "d", "e", "f", "g"}, - {0, 1, 2, 3, 4, 5, 6}, error); - require(invalid == nullptr, "oversized FAST horizons must fail before graph work"); - - std::vector> oversized_batch(1025); - require(!codec->decode_fast_tokens(oversized_batch, decoded, error), - "oversized FAST batches must fail before output allocation"); - - auto work_limited_codec = make_synthetic_codec(257, 1); - std::vector> excessive_idct_batch(1024); - require(!work_limited_codec->decode_fast_tokens(excessive_idct_batch, decoded, error) && - error.find("work limit") != std::string::npos, - "inverse-DCT work accounting must include the full batch dimension"); -} - -} // namespace - -int main() { - run_unit_tests(); - std::cout << "starvla FAST codec unit tests passed\n"; - return 0; -} diff --git a/tests/starvla/fast_runtime_test.cpp b/tests/starvla/fast_runtime_test.cpp deleted file mode 100644 index 7ec1dc7..0000000 --- a/tests/starvla/fast_runtime_test.cpp +++ /dev/null @@ -1,137 +0,0 @@ -#include "models/starvla/fast_codec.h" -#include "models/starvla/fast_policy.h" -#include "models/starvla/qwen3vl_bridge.h" - -#include -#include -#include -#include -#include -#include -#include -#include - -namespace { - -using robotcpp::starvla::FastCodec; -using robotcpp::starvla::FastCodecConfig; -using robotcpp::starvla::FastDecodeResult; -using robotcpp::starvla::FastPolicy; - -void require(bool condition, const std::string & message) { - if (!condition) { - std::cerr << "FAIL: " << message << '\n'; - std::exit(1); - } -} - -void test_generation_selector() { - std::string error; - int32_t token = -1; - std::vector logits = {10.0f, 9.0f, 1.0f}; - require(robotcpp::starvla::qwen_vl_select_repetition_penalized_top1( - logits.data(), logits.size(), {0}, 2.0f, token, error), - "repetition-penalized selector must succeed: " + error); - require(token == 1, - "a repeated positive logit must be divided before top_k=1"); - - logits = {-1.0f, -1.5f, -4.0f}; - require(robotcpp::starvla::qwen_vl_select_repetition_penalized_top1( - logits.data(), logits.size(), {0}, 2.0f, token, error), - "negative-logit selector must succeed"); - require(token == 1, - "a repeated negative logit must be multiplied before top_k=1"); - - logits = {3.0f, 3.0f}; - require(robotcpp::starvla::qwen_vl_select_repetition_penalized_top1( - logits.data(), logits.size(), {}, 1.05f, token, error) && - token == 0, - "top_k=1 tie handling must match torch.argmax first-index semantics"); - require(!robotcpp::starvla::qwen_vl_select_repetition_penalized_top1( - logits.data(), logits.size(), {2}, 1.05f, token, error), - "out-of-vocabulary history must be rejected"); - logits[0] = std::numeric_limits::quiet_NaN(); - require(!robotcpp::starvla::qwen_vl_select_repetition_penalized_top1( - logits.data(), logits.size(), {}, 1.05f, token, error), - "NaN generation logits must fail closed"); -} - -void test_compiled_codec() { - FastCodecConfig config; - config.scale = 1.0; - config.min_token = 0; - config.vocab_size = 4; - config.time_horizon = 2; - config.action_dim = 2; - std::string error; - auto codec = FastCodec::create_compiled( - config, {0, 1, 2, 3, 4}, {10, 20, 30, 40}, - {100, 101, 102, 103}, error); - require(codec != nullptr, - "compiled FAST codec must construct without sidecars: " + error); - - FastDecodeResult decoded; - require(codec->decode_generated_tokens( - {{999, 100, 101, 888, 102, 103}}, decoded, error), - "compiled FAST codec must filter and decode a full Qwen sequence"); - const double root_half = std::sqrt(0.5); - const std::vector expected = { - root_half * 40.0, root_half * 60.0, - root_half * -20.0, root_half * -20.0, - }; - require(decoded.actions.size() == expected.size(), - "compiled FAST codec must return the configured 2x2 action shape"); - for (size_t i = 0; i < expected.size(); ++i) { - require(std::fabs(decoded.actions[i] - expected[i]) < 1.0e-12, - "compiled FAST codec IDCT differs from the analytical result"); - } - - require(FastCodec::create_compiled( - config, {0, 1, 1, 3, 4}, {10, 20, 30, 40}, - {100, 101, 102, 103}, error) == nullptr, - "compiled FAST codec must reject non-increasing offsets"); -} - -void test_policy(const std::string & path) { - std::string error; - std::unique_ptr policy = FastPolicy::load(path, 0, error); - require(policy != nullptr, "official FAST policy GGUF must load: " + error); - const auto & config = policy->config(); - require(config.bundle_uuid == "b2651406-918b-524b-9df6-66861d744f29" && - config.qwen_hidden_dim == 2048 && - config.qwen_vocab_size == 153713 && - config.generation_max_length == 2048 && - config.generation_eos_token_ids == - std::vector({151645, 151643}), - "official FAST policy metadata must expose the pinned runtime"); - - std::vector action_ids; - std::vector fast_ids; - std::vector normalized; - require(!policy->decode_generated( - {100, 151665, 200}, action_ids, fast_ids, normalized, error), - "incomplete FAST output must fail"); - require(action_ids == std::vector({151665}) && - fast_ids == std::vector({0}) && normalized.empty(), - "failed FAST decode must not return actions"); -} - -} // namespace - -int main(int argc, char ** argv) { - test_generation_selector(); - test_compiled_codec(); - - if (argc == 3 && std::string(argv[1]) == "--policy") { - test_policy(argv[2]); - } else if (argc == 3 && std::string(argv[1]) == "--expect-reject") { - std::string error; - require(FastPolicy::load(argv[2], 0, error) == nullptr && !error.empty(), - "tampered FAST policy GGUF must fail closed"); - } else if (argc != 1) { - std::cerr << "usage: " << argv[0] - << " [--policy|--expect-reject ]\n"; - return 2; - } - return 0; -} diff --git a/tests/starvla/groot_prompt_test.cpp b/tests/starvla/groot_prompt_test.cpp deleted file mode 100644 index fe06d0d..0000000 --- a/tests/starvla/groot_prompt_test.cpp +++ /dev/null @@ -1,86 +0,0 @@ -#include "models/starvla/groot_prompt.h" - -#include -#include -#include - -namespace { - -void require(bool condition, const char * message) { - if (!condition) { - std::cerr << "FAIL: " << message << '\n'; - std::exit(1); - } -} - -void require_rejected(const std::string & cot_template, const std::string & task, - const char * message) { - std::string instruction = "stale instruction"; - std::string error = "stale error"; - require(!robotcpp::starvla::build_groot_instruction(cot_template, task, instruction, error), - message); - require(instruction.empty(), "rejected prompts must clear their instruction output"); - require(!error.empty(), "rejected prompts must explain the contract violation"); -} - -} // namespace - -int main() { - using robotcpp::starvla::build_groot_instruction; - using robotcpp::starvla::build_pi_v3_instruction; - - const std::string official_template = - "Your task is {instruction}. To identify the key objects for your task. " - "Locate their bounding boxes in [x1,y1,x2,y2] format."; - - std::string instruction; - std::string error = "stale error"; - require(build_groot_instruction(official_template, "grab the block", instruction, error), - "official GR00T prompt must build"); - require(instruction == - "Your task is grab the block. To identify the key objects for your task. " - "Locate their bounding boxes in [x1,y1,x2,y2] format.", - "official GR00T prompt must preserve the checkpoint template exactly"); - require(error.empty(), "successful prompt construction must clear stale errors"); - - require(build_pi_v3_instruction(official_template, "grab the block", instruction, error), - "official PI_v3 prompt must build"); - require(instruction == - "Your task is grab the block. To identify the key objects for your task. " - "Locate their bounding boxes in [x1,y1,x2,y2] format.", - "PI_v3 must use the same pinned CoT replacement contract"); - require(!build_pi_v3_instruction(official_template, "", instruction, error), - "empty PI_v3 tasks must be rejected"); - require(error.find("PI_v3") != std::string::npos, - "PI_v3 prompt errors must identify the active framework"); - - require(build_groot_instruction(official_template, "grab the block.", instruction, error), - "punctuated tasks must build"); - require(instruction == - "Your task is grab the block.. To identify the key objects for your task. " - "Locate their bounding boxes in [x1,y1,x2,y2] format.", - "task punctuation must not be normalized"); - - require(build_groot_instruction("First {instruction}; then {instruction}.", "pick", instruction, - error), - "templates with repeated placeholders must build"); - require(instruction == "First pick; then pick.", - "every instruction placeholder must be replaced"); - - require_rejected(official_template, "", "empty tasks must be rejected"); - require_rejected(official_template, std::string("grab\0now", 8), - "embedded NUL bytes in tasks must be rejected"); - require_rejected(std::string("Use {instruction}\0now", 21), "grab", - "embedded NUL bytes in templates must be rejected"); - require_rejected(official_template, "grab <__media__> now", - "reserved mtmd media markers in tasks must be rejected"); - require_rejected("<__media__>{instruction}", "grab", - "reserved mtmd media markers in templates must be rejected"); - require_rejected("Your task is ready.", "grab", - "templates without an instruction placeholder must be rejected"); - require_rejected("Your task is { instruction }.", "grab", - "lookalike placeholders must not satisfy the template contract"); - - std::cout << "starvla GR00T prompt tests passed\n"; - return 0; -} diff --git a/tests/starvla/model_test.cpp b/tests/starvla/model_test.cpp new file mode 100644 index 0000000..73f665b --- /dev/null +++ b/tests/starvla/model_test.cpp @@ -0,0 +1,59 @@ +#include "models/model.h" +#include "models/starvla/starvla_engine.h" + +#include +#include +#include +#include + +namespace { + +using robotcpp::starvla::StarVLAVariant; + +struct VariantCase { + const char * framework; + const char * backbone; + const char * name; + StarVLAVariant variant; +}; + +constexpr std::array kVariants = {{ + // Qwen3-VL + {"oft", "qwen3_vl", "qwen3_oft", StarVLAVariant::qwen3_oft}, + {"groot", "qwen3_vl", "qwen3_groot", StarVLAVariant::qwen3_groot}, + {"pi_v3", "qwen3_vl", "qwen3_pi_v3", StarVLAVariant::qwen3_pi_v3}, + + // Qwen2.5-VL + {"oft", "qwen2_5_vl", "qwen25_oft", StarVLAVariant::qwen25_oft}, + {"groot", "qwen2_5_vl", "qwen25_groot", StarVLAVariant::qwen25_groot}, + {"pi", "qwen2_5_vl", "qwen25_pi", StarVLAVariant::qwen25_pi}, + {"fast", "qwen2_5_vl", "qwen25_fast", StarVLAVariant::qwen25_fast}, +}}; + +} // namespace + +int main() { + for (const VariantCase & test : kVariants) { + StarVLAVariant variant = StarVLAVariant::qwen3_oft; + if (!robotcpp::starvla::starvla_variant_from_metadata( + test.framework, test.backbone, variant) || + variant != test.variant || + std::string(robotcpp::starvla::starvla_variant_name(variant)) != test.name || + std::string(robotcpp::starvla::starvla_variant_framework(variant)) != + test.framework) { + std::fprintf(stderr, "variant check failed: %s\n", test.name); + return 1; + } + } + + robotcpp::model_args args; + args.type = robotcpp::model_type::starvla; + std::unique_ptr model; + std::string error; + if (robotcpp::make_model(args, model, error) || model || + error.find("policy path is required") == std::string::npos) { + std::fprintf(stderr, "unexpected factory result: %s\n", error.c_str()); + return 1; + } + return 0; +} diff --git a/tests/starvla/oft_image_preprocess_test.cpp b/tests/starvla/oft_image_preprocess_test.cpp deleted file mode 100644 index 89f7d95..0000000 --- a/tests/starvla/oft_image_preprocess_test.cpp +++ /dev/null @@ -1,162 +0,0 @@ -#include "models/starvla/oft_image_preprocess.h" - -#include -#include -#include -#include -#include - -namespace { - -void require(bool condition, const char * message) { - if (!condition) { - std::cerr << "FAIL: " << message << '\n'; - std::exit(1); - } -} - -uint64_t fnv1a(const std::vector & values) { - uint64_t hash = UINT64_C(14695981039346656037); - for (uint8_t value : values) { - hash ^= value; - hash *= UINT64_C(1099511628211); - } - return hash; -} - -std::vector test_image(int width, int height, int stride) { - std::vector image(static_cast(stride) * height, 0xee); - for (int y = 0; y < height; ++y) { - for (int x = 0; x < width; ++x) { - for (int c = 0; c < 3; ++c) { - image[static_cast(y) * stride + x * 3 + c] = - static_cast((x * 37 + y * 61 + c * 83 + x * y * 7) & 0xff); - } - } - } - return image; -} - -} // namespace - -int main() { - using namespace robotcpp::starvla; - - const int width = 7; - const int height = 5; - const int stride = width * 3 + 5; - const std::vector source = test_image(width, height, stride); - std::vector pillow; - std::vector torch; - std::string error; - require(resize_pillow_bicubic_rgb(source.data(), width, height, stride, 4, 6, pillow, error), - "Pillow bicubic resize must succeed"); - require(resize_torchvision_bicubic_aa_rgb(source.data(), width, height, stride, 8, 9, torch, error), - "torchvision bicubic resize must succeed"); - require(pillow.size() == 4 * 6 * 3, "Pillow resize shape must be exact"); - require(torch.size() == 8 * 9 * 3, "torchvision resize shape must be exact"); - - // Byte-for-byte references: Pillow 7.0.0 and torchvision's uint8 bicubic-AA path. - require(fnv1a(pillow) == UINT64_C(13888783175115780895), - "Pillow bicubic resize must match the reference bytes"); - require(fnv1a(torch) == UINT64_C(14284510537394890340), - "torchvision bicubic-AA resize must match the reference bytes"); - - std::vector processed; - require(preprocess_oft_rgb(source.data(), width, height, 3, stride, 4, 6, 8, 9, - processed, error), - "two-stage OFT preprocessing must succeed"); - require(processed.size() == 8 * 9 * 3, "two-stage OFT image shape must be exact"); - require(fnv1a(processed) == UINT64_C(11011422699469149164), - "two-stage OFT preprocessing must match the reference bytes"); - - int smart_width = 0; - int smart_height = 0; - require(qwen3vl_smart_resize_dimensions(640, 488, 32, 65536, 16777216, - smart_width, smart_height, error), - "Qwen3-VL smart resize dimensions must succeed"); - require(smart_width == 640 && smart_height == 480, - "640x488 must map to the official 640x480 grid"); - require(qwen3vl_smart_resize_dimensions(640, 400, 32, 65536, 16777216, - smart_width, smart_height, error), - "Qwen3-VL ties-to-even smart resize must succeed"); - require(smart_width == 640 && smart_height == 384, - "Python ties-to-even rounding must map 400 pixels to 384"); - - std::vector smart; - int smart_tokens = 0; - require(preprocess_qwen3vl_rgb(source.data(), width, height, 3, stride, 16, 2, - 65536, 16777216, smart, smart_width, smart_height, - smart_tokens, error), - "Qwen3-VL dynamic preprocessing must succeed"); - require(smart_width == 320 && smart_height == 224 && smart_tokens == 70, - "Qwen3-VL dynamic preprocessing must report its exact grid and token count"); - require(smart.size() == static_cast(smart_width) * smart_height * 3, - "Qwen3-VL dynamic image shape must be exact"); - - const int official_width = 640; - const int official_height = 488; - const int official_stride = official_width * 3 + 5; - const std::vector official_shape_source = - test_image(official_width, official_height, official_stride); - require(preprocess_qwen3vl_rgb( - official_shape_source.data(), official_width, official_height, 3, - official_stride, 16, 2, 65536, 16777216, smart, smart_width, - smart_height, smart_tokens, error), - "official-shape Qwen3-VL preprocessing must succeed"); - require(smart_width == 640 && smart_height == 480 && smart_tokens == 300, - "official-shape Qwen3-VL preprocessing must produce a 30x40 merged grid"); - // Reference generated by torchvision 0.21's uint8 bicubic antialias path. - require(fnv1a(smart) == UINT64_C(11336294493867056015), - "official-shape Qwen3-VL resize must match torchvision bytes"); - - require(preprocess_qwen3vl_rgb( - official_shape_source.data(), official_width, official_height, 3, - official_stride, 14, 2, 3136, 12845056, smart, smart_width, - smart_height, smart_tokens, error), - "official-shape Qwen2.5-VL preprocessing must succeed"); - require(smart_width == 644 && smart_height == 476 && smart_tokens == 391, - "Qwen2.5-VL preprocessing must produce its exact 17x23 merged grid"); - // Reference generated by Transformers 4.57 / torchvision 0.21 fast Qwen2-VL - // image preprocessing on the same uint8 RGB tensor. - require(fnv1a(smart) == UINT64_C(12050652900109057577), - "Qwen2.5-VL resize must match torchvision bytes"); - - const int bridge_width = 256; - const int bridge_height = 256; - const int bridge_stride = bridge_width * 3; - const std::vector bridge_source = - test_image(bridge_width, bridge_height, bridge_stride); - require(preprocess_qwen3vl_rgb( - bridge_source.data(), bridge_width, bridge_height, 3, - bridge_stride, 14, 2, 3136, 12845056, smart, smart_width, - smart_height, smart_tokens, error), - "Bridge-size Qwen2.5-VL preprocessing must succeed"); - require(smart_width == 252 && smart_height == 252 && smart_tokens == 81, - "256x256 Bridge input must map to the official 18x18 patch grid"); - require(smart.size() == static_cast(252 * 252 * 3), - "Bridge-size Qwen2.5-VL preprocessing must return 252x252 RGB"); - - const int adapter_width = 224; - const int adapter_height = 224; - const int adapter_stride = adapter_width * 3; - const std::vector adapter_source = - test_image(adapter_width, adapter_height, adapter_stride); - require(preprocess_qwen3vl_rgb( - adapter_source.data(), adapter_width, adapter_height, 3, - adapter_stride, 14, 2, 3136, 12845056, smart, smart_width, - smart_height, smart_tokens, error), - "SimplerEnv adapter-size Qwen2.5-VL preprocessing must succeed"); - require(smart_width == 224 && smart_height == 224 && smart_tokens == 64, - "224x224 SimplerEnv input must preserve its official 16x16 patch grid"); - require(smart.size() == static_cast(224 * 224 * 3), - "SimplerEnv adapter-size preprocessing must return 224x224 RGB"); - - require(!qwen3vl_smart_resize_dimensions(640, 3, 32, 65536, 16777216, - smart_width, smart_height, error), - "Qwen3-VL smart resize must reject aspect ratios over 200"); - - std::cout << "pillow_fnv=" << fnv1a(pillow) << " torch_fnv=" << fnv1a(torch) - << " two_stage_fnv=" << fnv1a(processed) << '\n'; - return 0; -} diff --git a/tests/starvla/oft_prompt_test.cpp b/tests/starvla/oft_prompt_test.cpp deleted file mode 100644 index 5c6b838..0000000 --- a/tests/starvla/oft_prompt_test.cpp +++ /dev/null @@ -1,81 +0,0 @@ -#include "models/starvla/oft_prompt.h" - -#include -#include -#include -#include - -namespace { - -void require(bool condition, const char * message) { - if (!condition) { - std::cerr << "FAIL: " << message << '\n'; - std::exit(1); - } -} - -robotcpp::starvla::OFTPromptConfig official_config() { - robotcpp::starvla::OFTPromptConfig config; - config.horizon = 16; - config.action_token = "\xF0\x9F\x94\x8D"; - config.action_suffix = " Please predict the next 16 robot actions: "; - for (int i = 0; i < config.horizon; ++i) { - config.action_suffix += config.action_token; - } - config.action_suffix += "."; - config.cot_enabled = true; - config.cot_template = "Your task is {instruction}. Locate {instruction}."; - config.state_bins = 256; - config.state_bin_min = -1.0f; - config.state_bin_max = 1.0f; - config.state_clip = false; - return config; -} - -} // namespace - -int main() { - using namespace robotcpp::starvla; - - OFTPromptConfig config = official_config(); - std::string error; - require(validate_oft_prompt_config(config, error), "official prompt config must validate"); - - config.cot_enabled = false; - std::string instruction; - require(build_oft_instruction(config, "grab", {}, instruction, error), - "prompt without state must build"); - require(instruction == "grab" + config.action_suffix, "action suffix placement must match StarVLA"); - require(!build_oft_instruction(config, "grab <__media__> now", {}, instruction, error), - "reserved mtmd media markers in tasks must be rejected"); - require(error.find("reserved mtmd media marker") != std::string::npos, - "media marker rejection must explain the contract violation"); - - require(build_oft_instruction(config, "grab", {-1.1f, -1.0f, 0.0f, 1.0f, 1.1f}, - instruction, error), - "state prompt must build"); - require(instruction == "grab [STATE] -1 0 128 255 255 [ACTION]" + config.action_suffix, - "state discretization must match numpy.digitize"); - - config.cot_enabled = true; - require(build_oft_instruction(config, "grab", {}, instruction, error), "CoT prompt must build"); - const std::string unwrapped = "grab" + config.action_suffix; - require(instruction == "Your task is " + unwrapped + ". Locate " + unwrapped + ".", - "Python str.replace semantics must replace every instruction placeholder"); - - const std::string content = build_qwen_media_content(2, instruction, "<__media__>"); - require(content == "<__media__><__media__>" + instruction, - "image markers must precede text without separators"); - - std::vector ids = {1, 9, 2, 9, 9, 3, 9}; - std::vector positions; - require(find_last_token_positions(ids, 9, 3, positions, error), "action tokens must be found"); - require(positions == std::vector({3, 4, 6}), - "last action positions must remain in temporal order"); - require(!find_last_token_positions(ids, 9, 5, positions, error), - "insufficient action tokens must fail"); - require(!error.empty(), "insufficient action tokens must return an error"); - - std::cout << "starvla OFT prompt tests passed\n"; - return 0; -} diff --git a/tests/starvla/test_starvla_pi_v3_golden.py b/tests/starvla/test_starvla_pi_v3_golden.py deleted file mode 100644 index 4551bbd..0000000 --- a/tests/starvla/test_starvla_pi_v3_golden.py +++ /dev/null @@ -1,36 +0,0 @@ -from __future__ import annotations - -import sys -import unittest -from pathlib import Path - - -TOOLS_DIR = Path(__file__).resolve().parents[2] / "tools" / "hf2gguf" / "starvla" -sys.path.insert(0, str(TOOLS_DIR)) - -from generate_starvla_pi_v3_golden import ( # noqa: E402 - CONDITIONING_TAP_NAMES, - expected_model_instruction, - expected_runtime_contract, -) - - -class PIv3ReferenceTest(unittest.TestCase): - def test_instruction_template(self) -> None: - config = { - "datasets": { - "vla_data": {"CoT_prompt": "Task: {instruction}"}, - } - } - self.assertEqual(expected_model_instruction(config, "grab block"), "Task: grab block") - - def test_action_oracle_contract(self) -> None: - contract = expected_runtime_contract() - self.assertEqual(contract["conditioning"]["hidden_tuple_indices"], list(range(1, 37))) - self.assertEqual(contract["conditioning"]["hidden_tap_names"], CONDITIONING_TAP_NAMES) - self.assertEqual(contract["timesteps"], [0, 250, 500, 750]) - self.assertEqual(contract["action_shape"], [16, 7]) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/starvla/test_starvla_qwen25_fast.py b/tests/starvla/test_starvla_qwen25_fast.py deleted file mode 100644 index a8c7486..0000000 --- a/tests/starvla/test_starvla_qwen25_fast.py +++ /dev/null @@ -1,111 +0,0 @@ -from __future__ import annotations - -import sys -import tempfile -import unittest -from pathlib import Path - -import numpy as np - - -REPO_ROOT = Path(__file__).resolve().parents[2] -TOOLS_DIR = REPO_ROOT / "tools" / "hf2gguf" / "starvla" -sys.path.insert(0, str(TOOLS_DIR)) - -import convert_starvla_qwen25_fast as converter # noqa: E402 -from starvla_checkpoint import ( # noqa: E402 - StarVLAError, - load_catalog, - official_bundle_uuid, -) - - -SOURCE_ROOT = REPO_ROOT / "ckpts" / "starvla" / "sources" -POLICY_DIR = SOURCE_ROOT / "qwen25-fast-bridge-rt1" / "d9e2977d21755e78a0dd5f9a61586075a636d669" -QWEN_DIR = SOURCE_ROOT / "qwen2.5-vl-3b-instruct-action" / "ce86bd9a53416527b8361e8dfc47316288ffa110" -CODEC_DIR = SOURCE_ROOT / "fast-codec" / "ec4d7aa71691cac0b8bed6942be45684db2110f4" - - -def runtime_inputs() -> tuple[dict[str, object], dict[str, object], dict[str, object]]: - catalog = load_catalog() - entry, qwen, codec = converter.validate_catalog_contract(catalog) - manifest = { - "bundle_uuid": official_bundle_uuid(entry, catalog), - "source": { - "starvla_revision": catalog["source_revisions"]["starvla"], - "llama_cpp_revision": catalog["source_revisions"]["llama_cpp"], - "qwen_repo_id": qwen["repo_id"], - "qwen_revision": qwen["revision"], - }, - } - return manifest, entry, codec - - -@unittest.skipUnless( - POLICY_DIR.is_dir() and QWEN_DIR.is_dir() and CODEC_DIR.is_dir(), - "pinned FAST assets are not available", -) -class Qwen25FastTest(unittest.TestCase): - def test_preflight_and_codec_tables(self) -> None: - report = converter.preflight(load_catalog(), POLICY_DIR, QWEN_DIR, CODEC_DIR) - arrays = converter.compile_fast_runtime_tensors(QWEN_DIR, CODEC_DIR) - - self.assertEqual(report["variant"], "qwen25_fast") - self.assertEqual(set(arrays), converter.FAST_RUNTIME_TENSOR_NAMES) - self.assertEqual( - arrays[converter.ACTION_TOKEN_MAP_TENSOR].shape, - (converter.ACTION_TOKEN_COUNT,), - ) - self.assertEqual( - arrays[converter.CODEC_TOKEN_OFFSETS_TENSOR].shape, - (converter.ACTION_TOKEN_COUNT + 1,), - ) - - def test_runtime_policy_round_trip(self) -> None: - manifest, entry, codec = runtime_inputs() - metadata, arrays = converter.build_fast_runtime_policy( - manifest=manifest, - entry=entry, - codec_entry=codec, - source_dir=POLICY_DIR, - qwen_dir=QWEN_DIR, - codec_dir=CODEC_DIR, - ) - self.assertEqual(metadata["starvla.framework"], "fast") - self.assertEqual(metadata["starvla.model_type"], "starvla") - self.assertNotIn("starvla.fast.codec.decode_fallback", metadata) - self.assertNotIn("starvla.fast.runtime_contract_json", metadata) - - with tempfile.TemporaryDirectory() as temporary: - path = Path(temporary) / "policy.gguf" - converter.write_fast_runtime_policy_gguf(path, metadata, arrays) - record = converter.validate_fast_runtime_policy_gguf( - path, - expected_metadata=metadata, - expected_arrays=arrays, - ) - self.assertEqual(record["tensor_count"], 3) - - def test_validator_rejects_duplicate_action_token_ids(self) -> None: - manifest, entry, codec = runtime_inputs() - metadata, arrays = converter.build_fast_runtime_policy( - manifest=manifest, - entry=entry, - codec_entry=codec, - source_dir=POLICY_DIR, - qwen_dir=QWEN_DIR, - codec_dir=CODEC_DIR, - ) - arrays = dict(arrays) - action_map = np.array(arrays[converter.ACTION_TOKEN_MAP_TENSOR], copy=True) - action_map[1] = action_map[0] - arrays[converter.ACTION_TOKEN_MAP_TENSOR] = action_map - with tempfile.TemporaryDirectory() as temporary: - path = Path(temporary) / "policy.gguf" - converter.write_fast_runtime_policy_gguf(path, metadata, arrays) - with self.assertRaisesRegex(StarVLAError, "codec tensors"): - converter.validate_fast_runtime_policy_gguf(path) - - -if __name__ == "__main__": - unittest.main() diff --git a/tools/llama_cpp/apply_starvla_patches.sh b/tools/apply_patches.sh similarity index 91% rename from tools/llama_cpp/apply_starvla_patches.sh rename to tools/apply_patches.sh index be471c1..746584c 100755 --- a/tools/llama_cpp/apply_starvla_patches.sh +++ b/tools/apply_patches.sh @@ -13,10 +13,10 @@ readonly PATCHES=( usage() { cat <<'EOF' -Usage: tools/llama_cpp/apply_starvla_patches.sh [--check|--revert] +Usage: tools/apply_patches.sh [--check|--revert] -With no option, apply the StarVLA patches to third_party/llama.cpp. - --check Validate the pinned revision and report patch state. +With no option, apply the repository patches to third-party checkouts. + --check Validate pinned revisions and report patch state. --revert Remove an already applied complete patch set. Set LLAMA_CPP_DIR to validate or patch another checkout of the pinned revision. @@ -77,7 +77,7 @@ fi if [[ "${mode}" == "apply" ]]; then if ${all_applied}; then - echo "StarVLA llama.cpp patches are already applied." + echo "Repository patches are already applied." exit 0 fi if ! ${all_pending}; then @@ -99,7 +99,7 @@ if [[ "${mode}" == "apply" ]]; then fi if ${all_pending}; then - echo "StarVLA llama.cpp patches are not applied." + echo "Repository patches are not applied." exit 0 fi if ! ${all_applied}; then diff --git a/tools/hf2gguf/README.md b/tools/hf2gguf/README.md index 3cf2fc0..2f4260c 100644 --- a/tools/hf2gguf/README.md +++ b/tools/hf2gguf/README.md @@ -8,7 +8,8 @@ This directory contains tools for converting checkpoints to GGUF. - `smolvla/`: converts LeRobot-style SmolVLA checkpoints into four GGUF components. - `pi0/`: converts LeRobot-style pi0 checkpoints into six split GGUF components. -- `starvla/`: pinned Qwen3-VL and Qwen2.5-VL conversion plus a shared 3% CUDA action parity gate for seven StarVLA variants. No official finetuned Qwen3 FAST policy checkpoint is available. +- `starvla/`: converts seven StarVLA Qwen3-VL and Qwen2.5-VL checkpoints. See + [`starvla/README.md`](starvla/README.md) for supported variants and commands. - `environment.yaml`: conda environment for the converters. ## Usage diff --git a/tools/hf2gguf/README_ZH.md b/tools/hf2gguf/README_ZH.md index f0fb6e4..ad1cec1 100644 --- a/tools/hf2gguf/README_ZH.md +++ b/tools/hf2gguf/README_ZH.md @@ -8,7 +8,8 @@ - `smolvla/`：将 SmolVLA的lerobot-style的checkpoint 转成四个 GGUF component。 - `pi0/`：将 pi0的lerobot-style的checkpoint 转成六个 split GGUF component。 -- `starvla/`：固定版本的 Qwen3-VL/Qwen2.5-VL 转换与共享的 3% CUDA action parity gate，覆盖七种 StarVLA variant；上游尚无官方 finetuned Qwen3 FAST policy checkpoint。 +- `starvla/`：转换七个 StarVLA Qwen3-VL 和 Qwen2.5-VL checkpoint，支持范围和命令见 + [`starvla/README.md`](starvla/README.md)。 - `environment.yaml`：converter conda 环境。 ## 使用说明 diff --git a/tools/hf2gguf/starvla/README.md b/tools/hf2gguf/starvla/README.md new file mode 100644 index 0000000..a663e35 --- /dev/null +++ b/tools/hf2gguf/starvla/README.md @@ -0,0 +1,177 @@ +# Converting StarVLA checkpoints + +The scripts in this directory download StarVLA checkpoints and convert them to +GGUF files used by robot.cpp. + +## Models + +| Variant | Backbone | Policy | +| --- | --- | --- | +| `oft` | Qwen3-VL | OFT | +| `groot` | Qwen3-VL | GR00T | +| `pi_v3` | Qwen3-VL | PI_v3 | +| `qwen25_oft` | Qwen2.5-VL | OFT | +| `qwen25_groot` | Qwen2.5-VL | GR00T | +| `qwen25_pi` | Qwen2.5-VL | PI | +| `qwen25_fast` | Qwen2.5-VL | FAST | + +All variants use `starvla` as the public model type. The loader reads the +backbone and policy type from the policy GGUF. + +The first six variants contain a BF16 Qwen file, a BF16 multimodal projector, +and an FP32 policy file. FAST uses the fine-tuned BF16 Qwen model as its policy; +its separate policy GGUF contains the integer token map and codec data. The +loader checks the bundle UUID to prevent files from different conversions from +being combined. + +Qwen3 FAST is not listed because StarVLA has not published a fine-tuned Qwen3 +FAST policy checkpoint. + +[`checkpoint_catalog.json`](checkpoint_catalog.json) defines the supported +topologies and pins the official release files and shared Qwen assets. + +## Environment + +```bash +conda env create -f tools/hf2gguf/starvla/environment.yaml +conda activate starvla_gguf_converter +``` + +The scripts use `.venv/bin/python` by default. Set `PYTHON=python` to use the +active conda environment. + +## Convert an official release + +Convert one of the variants from the table above: + +```bash +tools/hf2gguf/starvla/convert.sh oft +``` + +This downloads and verifies the catalog checkpoint, prepares a +clean llama.cpp worktree at the pinned revision, converts all components, and +validates the resulting bundle. It refuses to overwrite an existing output +directory. Pass a second argument to select another output directory: + +```bash +tools/hf2gguf/starvla/convert.sh qwen25_fast /path/to/output +``` + +## Convert a training checkpoint + +Current StarVLA training runs contain the files needed by the converter: + +```text +/ + config.yaml + dataset_statistics.json + checkpoints/steps__pytorch_model.pt + # or checkpoints/steps__model.safetensors +``` + +Pass the checkpoint and the matching topology from the model table: + +```bash +tools/hf2gguf/starvla/convert.sh oft /path/to/output \ + --checkpoint /path/to/run/checkpoints/steps_5000_model.safetensors +``` + +The run directory is inferred from checkpoints under `checkpoints/` or +`final_model/`. Use `--source-dir /path/to/run` when the files use another +layout. If `dataset_statistics.json` has several profiles and does not contain +the catalog default, select one with `--unnorm-key`: + +```bash +tools/hf2gguf/starvla/convert.sh groot /path/to/output \ + --checkpoint /path/to/run/final_model/pytorch_model.pt \ + --unnorm-key bridge_dataset +``` + +Supported training exports are flat PyTorch state dictionaries (`.pt`) and +flat safetensors files (`.safetensors`) written by `train_starvla.py`, including +periodic and final checkpoints. The converter does not consume optimizer +state, distributed checkpoint shards, or a checkpoint whose architecture no +longer matches the selected variant. `config.json` and `config.full.yaml` are +not required. + +The converter hashes the local checkpoint and run metadata, so its bundle UUID +and manifest differ from the official release even when the weights are equal. + +A successful conversion writes exactly four files: + +```text +qwen--bf16.gguf +mmproj--bf16.gguf +starvla--policy-fp32.gguf +conversion_manifest.json +``` + +For FAST, the policy GGUF stores the integer token map and codec data instead +of FP32 policy weights. Its filenames are: + +```text +qwen-qwen25-fast-bf16.gguf +mmproj-qwen25-fast-bf16.gguf +policy-qwen25-fast.gguf +conversion_manifest.json +``` + +Set `STARVLA_LOCAL_FILES_ONLY=1` to forbid network access and use already +downloaded sources. The low-level converters remain available for debugging, +but normal conversion should use `convert.sh` so all paths and revisions come +from [`checkpoint_catalog.json`](checkpoint_catalog.json). + +## Build + +The runtime needs two llama.cpp patches maintained in this repository. See +[`patches/llama.cpp/README.md`](../../../patches/llama.cpp/README.md) for their +scope. + +```bash +./tools/apply_patches.sh +cmake -S . -B build_cuda \ + -DGGML_CUDA=ON \ + -DBUILD_TESTING=ON \ + -DROBOT_CPP_BUILD_STARVLA=ON \ + -DROBOT_CPP_BUILD_MODEL_CLI=ON +cmake --build build_cuda -j +``` + +## Run + +Pass the Qwen, multimodal projector, and policy GGUF files separately: + +```bash +CUDA_VISIBLE_DEVICES=0 build_cuda/bin/model-cli \ + --model-type starvla \ + --policy ckpts/starvla/gguf/oft/starvla-oft-policy-fp32.gguf \ + --llm ckpts/starvla/gguf/oft/qwen-oft-bf16.gguf \ + --mmproj ckpts/starvla/gguf/oft/mmproj-oft-bf16.gguf \ + --image /path/to/frame-224-rgb.png \ + --image-name image_0 \ + --task "grab the block." \ + --n-ctx 2048 \ + --n-batch 2048 +``` + +GR00T, PI_v3, and PI accept `--noise-seed`. FAST accepts one RGB `image_0` and +no robot state. + +The server uses the same model type and policy file: + +```bash +CUDA_VISIBLE_DEVICES=0 build_cuda/bin/model-server \ + --model-type starvla \ + --policy ckpts/starvla/gguf/oft/starvla-oft-policy-fp32.gguf \ + --llm ckpts/starvla/gguf/oft/qwen-oft-bf16.gguf \ + --mmproj ckpts/starvla/gguf/oft/mmproj-oft-bf16.gguf \ + --host 127.0.0.1 \ + --port 5555 \ + --n-ctx 2048 \ + --n-batch 2048 +``` + +The policy GGUF records its default action normalization profile. + +The repository does not include upstream checkpoints. Check each model's +license before distributing converted files. diff --git a/tools/hf2gguf/starvla/checkpoint_catalog.json b/tools/hf2gguf/starvla/checkpoint_catalog.json new file mode 100644 index 0000000..040c010 --- /dev/null +++ b/tools/hf2gguf/starvla/checkpoint_catalog.json @@ -0,0 +1,478 @@ +{ + "schema_version": 1, + "source_revisions": { + "starvla": "631aae02afe6d95876e923ff518e8ff2ab9a2f88", + "llama_cpp": "3e941b813b1acbbf06c2203a94ceb33d84748c1e" + }, + "shared_assets": { + "qwen3_vl_4b_instruct": { + "directory": "qwen3-vl-4b-instruct", + "repo_id": "Qwen/Qwen3-VL-4B-Instruct", + "revision": "ebb281ec70b05090aa6165b016eac8ec08e71b17", + "files": [ + "chat_template.json", + "config.json", + "generation_config.json", + "merges.txt", + "preprocessor_config.json", + "tokenizer.json", + "tokenizer_config.json", + "video_preprocessor_config.json", + "vocab.json" + ], + "file_hashes": { + "chat_template.json": {"size": 5502, "sha256": "6f8a6a55027e3da5160105556cda5dd69f6423f1c32645f6730d32de7773d0c4"}, + "config.json": {"size": 1505, "sha256": "edac7703329133edfc53e46ac0081835144c99d7eebf28b71c732694d435224d"}, + "generation_config.json": {"size": 269, "sha256": "8469742d1fce0de951c8909b26a2c0c0d8490837ce476efb114da9e0cefc4d44"}, + "merges.txt": {"size": 1671839, "sha256": "599bab54075088774b1733fde865d5bd747cbcc7a547c5bc12610e874e26f5e3"}, + "preprocessor_config.json": {"size": 390, "sha256": "27225450ac9c6529872ee1924fcb0962ff5634834f817040f444118116f4e516"}, + "tokenizer.json": {"size": 7032403, "sha256": "a5d85b6dcc535e6b93115a9ef287e6132fdbf30270da6218194ba742261173c7"}, + "tokenizer_config.json": {"size": 10868, "sha256": "c2da771801886ad9ae98181793ffd3dfb7f1af30f6f7c6a4e15d7dbba52e2399"}, + "video_preprocessor_config.json": {"size": 385, "sha256": "7768af27c1fafa9cc9011c1dc20067e03f8915e03b63504550e11d5066986d13"}, + "vocab.json": {"size": 2776833, "sha256": "ca10d7e9fb3ed18575dd1e277a2579c16d108e32f27439684afa0e10b1440910"} + }, + "staged_overrides": { + "config.json": {"size": 1507, "sha256": "ef6ec5fd4c5a80b549208f2352d88c480436db2cf9983359f23260c36e4ae38d"} + } + }, + "qwen2_5_vl_3b_instruct": { + "directory": "qwen2.5-vl-3b-instruct", + "repo_id": "Qwen/Qwen2.5-VL-3B-Instruct", + "revision": "66285546d2b821cf421d4f5eb2576359d3770cd3", + "files": [ + "chat_template.json", + "config.json", + "generation_config.json", + "merges.txt", + "model.safetensors.index.json", + "preprocessor_config.json", + "tokenizer.json", + "tokenizer_config.json", + "vocab.json" + ], + "file_hashes": { + "chat_template.json": {"size": 1050, "sha256": "ad60d90252ed0b0705ba14e2d0ad0fec0beac1ea955642b54059b36052d8bc96"}, + "config.json": {"size": 1373, "sha256": "7ed3eed5be6924cc800e8a5e53fc405c1aab1aaf36bad65c33403b36c56827f5"}, + "generation_config.json": {"size": 216, "sha256": "533f191cc257b7de37a4fccd0a7a1706d75e1aa660f93efaa54e5a2a9f9aace9"}, + "merges.txt": {"size": 1671839, "sha256": "599bab54075088774b1733fde865d5bd747cbcc7a547c5bc12610e874e26f5e3"}, + "model.safetensors.index.json": {"size": 65448, "sha256": "c7dd78a4c6bea60b51332f1baf37b8f8124ecab2c35395a29a29825bf2619768"}, + "preprocessor_config.json": {"size": 350, "sha256": "f2058c716eef96ccaed1cc1e2d0c08306b62586d535b28d9d08e691b2fab7ca0"}, + "tokenizer.json": {"size": 7031645, "sha256": "c0382117ea329cdf097041132f6d735924b697924d6f6fc3945713e96ce87539"}, + "tokenizer_config.json": {"size": 5702, "sha256": "4abd3520120e266da84c0864fee064d1fb10806f02225911a47253dd38dc5f56"}, + "vocab.json": {"size": 2776833, "sha256": "ca10d7e9fb3ed18575dd1e277a2579c16d108e32f27439684afa0e10b1440910"} + }, + "staged_overrides": { + "config.json": {"size": 1375, "sha256": "9c22fba5261a8e47aa66be0e4ef22473190168859dc3bbe7f283fbc4f161b0eb"} + }, + "optional_weight_files": [ + "model-00001-of-00002.safetensors", + "model-00002-of-00002.safetensors" + ], + "optional_weight_hashes": { + "model-00001-of-00002.safetensors": {"size": 3982649232, "sha256": "41a8895c164b4d32bae6b302f4603fcbc1797f32dafa45c7e9bcda23c6755df8"}, + "model-00002-of-00002.safetensors": {"size": 3526688744, "sha256": "365531ff8752420e89dee707b79d021fb2d6e25abafe486f080555a4fe6972e4"} + } + }, + "qwen2_5_vl_3b_instruct_action": { + "directory": "qwen2.5-vl-3b-instruct-action", + "repo_id": "StarVLA/Qwen2.5-VL-3B-Instruct-Action", + "revision": "ce86bd9a53416527b8361e8dfc47316288ffa110", + "files": [ + "README.md", + "added_token_id_map.json", + "added_tokens.json", + "chat_template.jinja", + "config.json", + "generation_config.json", + "merges.txt", + "model.safetensors.index.json", + "preprocessor_config.json", + "special_tokens_map.json", + "tokenizer.json", + "tokenizer_config.json", + "video_preprocessor_config.json", + "vocab.json" + ], + "file_hashes": { + "README.md": {"size": 482, "sha256": "cef1e9f3b90d50d1b6274fed603140127de9441bd6d8425811225d00a290f46f"}, + "added_token_id_map.json": {"size": 66476, "sha256": "a774a771870979578111a9f083e03e421bc3e6c0c7070d81e775acc21c74a21a"}, + "added_tokens.json": {"size": 67079, "sha256": "fcca65c62c6da071c4046abbc18b8287c51030e62febec408bac498a03652eaa"}, + "chat_template.jinja": {"size": 1017, "sha256": "a0bc6f6fc7a29a80017a433e8f03a1cc1236e838a944a2d034295a60c4f2fddb"}, + "config.json": {"size": 3317, "sha256": "5c30acf44442bbdd863b87a6f61b6879616a2933271bf62841de14037f6c0f7d"}, + "generation_config.json": {"size": 244, "sha256": "76001fd927297f839d96c5a52dd09de3406a4c28822fba4f525d13c5a1e2c8d7"}, + "merges.txt": {"size": 1671853, "sha256": "8831e4f1a044471340f7c0a83d7bd71306a5b867e95fd870f74d0c5308a904d5"}, + "model.safetensors.index.json": {"size": 65484, "sha256": "e6ce326cff552529deb7ca2e05616b9b79fba9d26632309564744c3731fbb644"}, + "preprocessor_config.json": {"size": 829, "sha256": "dfc7263fe735989c65c290d394198c4724d5afc58bf15c815d5e2e25b00b51b1"}, + "special_tokens_map.json": {"size": 312563, "sha256": "707f14d06c06e20212dbe5c118873f1c024f28ffac4b07c52bd2840ba0c34290"}, + "tokenizer.json": {"size": 11822194, "sha256": "07da2a694acc4f6e63d67da9926817ee35b0354b1e570a6a73d325760a1c2ed2"}, + "tokenizer_config.json": {"size": 438450, "sha256": "7cd59c7a865d2989c0d9b18bf485a5ee212f8ff333acd305a1bf560b75c16575"}, + "video_preprocessor_config.json": {"size": 913, "sha256": "15bb7c2f2bc95fe9cc3749a4b287872b4886a15e9fcf550a4122a46ae26150bd"}, + "vocab.json": {"size": 2776833, "sha256": "ca10d7e9fb3ed18575dd1e277a2579c16d108e32f27439684afa0e10b1440910"} + }, + "staged_overrides": { + "config.json": {"size": 3350, "sha256": "782edd73d2c9584d65350a6410780b96bef658437cbd9d8e0ed7006a1e3fcaed"} + }, + "optional_weight_files": [ + "model-00001-of-00002.safetensors", + "model-00002-of-00002.safetensors" + ], + "optional_weight_hashes": { + "model-00001-of-00002.safetensors": {"size": 4959940464, "sha256": "0abe459fc004959698441fd706b7721ec4633e963284b1b260df31ed5f765960"}, + "model-00002-of-00002.safetensors": {"size": 2556676080, "sha256": "a7a84f03ce697eddedc3abba904e2a110104fe78d6be9a0067f5b8f45c358e1c"} + } + }, + "fast_codec": { + "directory": "fast-codec", + "repo_id": "physical-intelligence/fast", + "revision": "ec4d7aa71691cac0b8bed6942be45684db2110f4", + "files": [ + "processing_action_tokenizer.py", + "processor_config.json", + "special_tokens_map.json", + "tokenizer.json", + "tokenizer_config.json" + ], + "file_hashes": { + "processing_action_tokenizer.py": {"size": 6145, "sha256": "6f021ca1f4c1b194ab6fa399d80baf3d642eadb17efb8f73301e4ac401522c20"}, + "processor_config.json": {"size": 253, "sha256": "f40cfbb1020858fe1d48c0f946b0c1315a90d6e84aa82710036f24f4c167706a"}, + "special_tokens_map.json": {"size": 3, "sha256": "ca3d163bab055381827226140568f3bef7eaac187cebd76878e0b63e9e442356"}, + "tokenizer.json": {"size": 686974, "sha256": "6507dd709287fd018882120c0071787f1f62bad9f180f1e8c5235bda1b71fa78"}, + "tokenizer_config.json": {"size": 322, "sha256": "b4030e2a13a0dea22e99d54c086fb320c71e66ad034ac4eba4301a0a27d5e5cd"} + } + } + }, + "variants": { + "oft": { + "model_type": "starvla", + "framework": "oft", + "backbone": "qwen3_vl", + "qwen_asset": "qwen3_vl_4b_instruct", + "default_unnorm_key": "oxe_bridge", + "directory": "oft-bridge-rt1", + "repo_id": "StarVLA/Qwen3VL-OFT-Bridge-RT-1", + "revision": "c3fc8f028429ba14819bf3b16e098776b670c889", + "files": [ + "config.json", + "config.yaml", + "dataset_statistics.json" + ], + "file_hashes": { + "config.json": {"size": 3920, "sha256": "6a6b0dd11ec26f88aca711a8886ecb619bbb895846a6353df53f04ece682b318"}, + "config.yaml": {"size": 3207, "sha256": "6c074974697115284b1624dda5230f3dd27e1d9a373db73006467004ed859c2a"}, + "dataset_statistics.json": {"size": 5990, "sha256": "83aa32682dd0b600e570936bfb63fd5d30b51d165e3d174a7ff3fc69d9fc276b"} + }, + "checkpoint": { + "path": "checkpoints/steps_5000_pytorch_model.pt", + "size": 9785060316, + "sha256": "371cb744227687bb99bcad7f9ff2250cf06da75631359ad3eba4c6bc52570607" + }, + "policy_prefixes": [ + "action_model." + ], + "expected": { + "total_tensors": 730, + "vlm_tensors": 714, + "policy_tensors": 16, + "visual_tensors": 315, + "text_tensors": 398, + "lm_head_tensors": 1, + "total_numel": 4892395015, + "vlm_numel": 4826771968, + "policy_numel": 65623047, + "dtypes": { + "bfloat16": 730 + }, + "storage_alias_groups": 0 + }, + "required_shapes": { + "model.language_model.embed_tokens.weight": [151936, 2560], + "lm_head.weight": [151936, 2560], + "action_model.model.fc1.weight": [5120, 2560], + "action_model.model.fc2.weight": [7, 5120] + } + }, + "groot": { + "model_type": "starvla", + "framework": "groot", + "backbone": "qwen3_vl", + "qwen_asset": "qwen3_vl_4b_instruct", + "default_unnorm_key": "oxe_bridge", + "directory": "groot-bridge-rt1", + "repo_id": "StarVLA/Qwen3VL-GR00T-Bridge-RT-1", + "revision": "12acc0b0f1f6230df21c479934a67a930b52f878", + "files": [ + "config.json", + "config.yaml", + "dataset_statistics.json" + ], + "file_hashes": { + "config.json": {"size": 3926, "sha256": "9efddc3c21039fa473823080a939dfa686050a8a2d4a4cb1b01b1a06913fccf5"}, + "config.yaml": {"size": 3174, "sha256": "01e092e9a3a9380885f1a27953048e7ed1ef7f6c99ea1984d289d1358e4ba85f"}, + "dataset_statistics.json": {"size": 5990, "sha256": "83aa32682dd0b600e570936bfb63fd5d30b51d165e3d174a7ff3fc69d9fc276b"} + }, + "checkpoint": { + "path": "checkpoints/steps_20000_pytorch_model.pt", + "size": 9976845210, + "sha256": "769d6c400d582a86ae8df8b0b445240ab679dbe77eeb72a4db71e43cd129c7c3" + }, + "policy_prefixes": [ + "action_model." + ], + "expected": { + "total_tensors": 962, + "vlm_tensors": 714, + "policy_tensors": 248, + "visual_tensors": 315, + "text_tensors": 398, + "lm_head_tensors": 1, + "total_numel": 4988244743, + "vlm_numel": 4826771968, + "policy_numel": 161472775, + "dtypes": { + "bfloat16": 962 + }, + "storage_alias_groups": 0 + }, + "required_shapes": { + "model.language_model.embed_tokens.weight": [151936, 2560], + "lm_head.weight": [151936, 2560] + } + }, + "pi_v3": { + "model_type": "starvla", + "framework": "pi_v3", + "backbone": "qwen3_vl", + "qwen_asset": "qwen3_vl_4b_instruct", + "default_unnorm_key": "oxe_bridge", + "directory": "pi-v3-bridge-rt1", + "repo_id": "StarVLA/Qwen3VL-PI_v3-Bridge-RT_1", + "revision": "99a3c01b3977e6442871a1fb62ce178279c5c3ed", + "files": [ + "config.full.yaml", + "config.yaml", + "dataset_statistics.json" + ], + "file_hashes": { + "config.full.yaml": {"size": 3087, "sha256": "85ff9fba2c9426d35f12efabfd52b9ef4633d91d084421952b391809ea73b33f"}, + "config.yaml": {"size": 1915, "sha256": "f266bd2de5b9fb7078c8314954e98f72476a36a8944e3f380096ba5b1256901b"}, + "dataset_statistics.json": {"size": 5987, "sha256": "9925e884e37ca807061b5d41206bcf72814300e7effb5b4edc067dc821fca004"} + }, + "checkpoint": { + "path": "checkpoints/steps_50000_pytorch_model.pt", + "size": 10922634912, + "sha256": "7f59a5d0fa9c167fabd941bca8e606bdf5597bfb4f99ca83e345672dd9c345ed" + }, + "policy_prefixes": [ + "action_model.", + "project_layers." + ], + "expected": { + "total_tensors": 1386, + "vlm_tensors": 714, + "policy_tensors": 672, + "visual_tensors": 315, + "text_tensors": 398, + "lm_head_tensors": 1, + "total_numel": 5461066247, + "vlm_numel": 4826771968, + "policy_numel": 634294279, + "dtypes": { + "bfloat16": 1386 + }, + "storage_alias_groups": 0 + }, + "required_shapes": { + "model.language_model.embed_tokens.weight": [151936, 2560], + "lm_head.weight": [151936, 2560] + } + }, + "qwen25_oft": { + "model_type": "starvla", + "framework": "oft", + "backbone": "qwen2_5_vl", + "qwen_asset": "qwen2_5_vl_3b_instruct", + "default_unnorm_key": "bridge_dataset", + "directory": "qwen25-oft-bridge-rt1", + "repo_id": "StarVLA/Qwen-OFT-Bridge-RT-1", + "revision": "11fa6440835ba3e912de43cfe8521043360ffc02", + "files": [ + "config.yaml", + "dataset_statistics.json", + "summary.jsonl" + ], + "file_hashes": { + "config.yaml": {"size": 2876, "sha256": "2f0362a1c0ae1eafba90d0feadf34652515c1d8a0c956db8f97e532492f2cdab"}, + "dataset_statistics.json": {"size": 6007, "sha256": "d2c4803c94d3b6eb1b8e8e100280e16a53b6058c8c2a7e747d27ecf9fcf9a4de"}, + "summary.jsonl": {"size": 33, "sha256": "a352646601877394e54b68ce09697098866d5904b825d7932826decaef0b2f8f"} + }, + "checkpoint": { + "path": "checkpoints/steps_10000_pytorch_model.pt", + "size": 8215912766, + "sha256": "51fe8d22c8d57116c2f59c5fdb24323fa3411149e888b807edba99b8354e0861" + }, + "policy_prefixes": [ + "action_model." + ], + "expected": { + "total_tensors": 841, + "vlm_tensors": 825, + "policy_tensors": 16, + "visual_tensors": 390, + "text_tensors": 434, + "lm_head_tensors": 1, + "total_numel": 4107800583, + "vlm_numel": 4065787904, + "policy_numel": 42012679, + "total_nbytes": 8215601166, + "vlm_nbytes": 8131575808, + "policy_nbytes": 84025358, + "dtypes": { + "bfloat16": 841 + }, + "storage_alias_groups": 0 + }, + "required_shapes": { + "model.embed_tokens.weight": [151936, 2048], + "lm_head.weight": [151936, 2048] + } + }, + "qwen25_groot": { + "model_type": "starvla", + "framework": "groot", + "backbone": "qwen2_5_vl", + "qwen_asset": "qwen2_5_vl_3b_instruct_action", + "default_unnorm_key": "oxe_bridge", + "directory": "qwen25-groot-bridge-rt1", + "repo_id": "StarVLA/Qwen-GR00T-Bridge-RT-1", + "revision": "5ebc661ba38b29c28f20fff6574801e6f49f3466", + "files": [ + "config.yaml", + "dataset_statistics.json", + "summary.jsonl" + ], + "file_hashes": { + "config.yaml": {"size": 3175, "sha256": "80d36dd087bd8d0feff246be94a7edcb296161823bce6340f76cc253724fbf1d"}, + "dataset_statistics.json": {"size": 5990, "sha256": "83aa32682dd0b600e570936bfb63fd5d30b51d165e3d174a7ff3fc69d9fc276b"}, + "summary.jsonl": {"size": 51, "sha256": "1fcbf58d35ac56b969410719240d507f91ce4ccbcf9c68540f30c8adf225b439"} + }, + "checkpoint": { + "path": "checkpoints/steps_30000_pytorch_model.pt", + "size": 8456891339, + "sha256": "9646da2ae0b32589a75c8cc88fae96c93c5d269b69fd7a29200744936e01d96f" + }, + "policy_prefixes": [ + "action_model." + ], + "expected": { + "total_tensors": 1073, + "vlm_tensors": 825, + "policy_tensors": 248, + "visual_tensors": 390, + "text_tensors": 434, + "lm_head_tensors": 1, + "total_numel": 4228247815, + "vlm_numel": 4073066496, + "policy_numel": 155181319, + "dtypes": { + "bfloat16": 1073 + }, + "storage_alias_groups": 0 + }, + "required_shapes": { + "model.embed_tokens.weight": [153713, 2048], + "lm_head.weight": [153713, 2048], + "action_model.model.timestep_encoder.timestep_embedder.linear_1.weight": [768, 256], + "action_model.model.transformer_blocks.0.attn1.to_k.weight": [768, 2048], + "action_model.model.transformer_blocks.1.attn1.to_k.weight": [768, 768], + "action_model.model.proj_out_2.weight": [1024, 768], + "action_model.action_decoder.layer2.weight": [7, 1024] + } + }, + "qwen25_pi": { + "model_type": "starvla", + "framework": "pi", + "backbone": "qwen2_5_vl", + "qwen_asset": "qwen2_5_vl_3b_instruct_action", + "default_unnorm_key": "oxe_bridge", + "directory": "qwen25-pi-bridge-rt1", + "repo_id": "StarVLA/Qwen-PI-Bridge-RT-1", + "revision": "26d0e079fbe3bc3fc62301f44f0025ef7c64ee22", + "files": [ + "config.yaml", + "dataset_statistics.json", + "summary.jsonl" + ], + "file_hashes": { + "config.yaml": {"size": 3154, "sha256": "a7bdbde311bc910ee81e673a899199035d372641562824570c4dbaba4ea99ee2"}, + "dataset_statistics.json": {"size": 5990, "sha256": "83aa32682dd0b600e570936bfb63fd5d30b51d165e3d174a7ff3fc69d9fc276b"}, + "summary.jsonl": {"size": 169, "sha256": "9051086fe35e01366062e6e3eee43c54b081af37f12f7b132446b22bef129d70"} + }, + "checkpoint": { + "path": "checkpoints/steps_30000_pytorch_model.pt", + "size": 10103104403, + "sha256": "8a0e47858921924d5038f7c4393dee6682b83175a85546e35e357e8f74ce8343" + }, + "policy_prefixes": [ + "action_model." + ], + "expected": { + "total_tensors": 1073, + "vlm_tensors": 825, + "policy_tensors": 248, + "visual_tensors": 390, + "text_tensors": 434, + "lm_head_tensors": 1, + "total_numel": 5051354119, + "vlm_numel": 4073066496, + "policy_numel": 978287623, + "total_nbytes": 10102708238, + "vlm_nbytes": 8146132992, + "policy_nbytes": 1956575246, + "dtypes": { + "bfloat16": 1073 + }, + "storage_alias_groups": 0 + }, + "required_shapes": { + "model.embed_tokens.weight": [153713, 2048], + "lm_head.weight": [153713, 2048], + "action_model.model.timestep_encoder.timestep_embedder.linear_1.weight": [2048, 256], + "action_model.model.transformer_blocks.0.attn1.to_k.weight": [2048, 2048], + "action_model.model.transformer_blocks.15.attn1.to_k.weight": [2048, 2048], + "action_model.model.proj_out_2.weight": [1024, 2048], + "action_model.action_decoder.layer2.weight": [7, 2048] + } + }, + "qwen25_fast": { + "model_type": "starvla", + "framework": "fast", + "backbone": "qwen2_5_vl", + "qwen_asset": "qwen2_5_vl_3b_instruct_action", + "default_unnorm_key": "bridge_dataset", + "directory": "qwen25-fast-bridge-rt1", + "repo_id": "StarVLA/Qwen-FAST-Bridge-RT-1", + "revision": "d9e2977d21755e78a0dd5f9a61586075a636d669", + "files": [ + "config.yaml", + "dataset_statistics.json", + "summary.jsonl" + ], + "file_hashes": { + "config.yaml": {"size": 2841, "sha256": "c0520794c8e5a15841b09fc1d9fb43216394674510d60e1f6d27dec149fa34f2"}, + "dataset_statistics.json": {"size": 6007, "sha256": "d2c4803c94d3b6eb1b8e8e100280e16a53b6058c8c2a7e747d27ecf9fcf9a4de"}, + "summary.jsonl": {"size": 34, "sha256": "2e927f0dd6524ec7cad6bb3023153142ee1c1f8015ba3712fcb0581dbf8c95e0"} + }, + "checkpoint": { + "path": "checkpoints/steps_10000_pytorch_model.pt", + "size": 8146439050, + "sha256": "f30e89a6b2a166fa3f48af42d5cffde07be44074b861abc7b57e1ccdb734e81e" + }, + "policy_prefixes": [], + "expected": null, + "required_shapes": { + "model.embed_tokens.weight": [153713, 2048], + "lm_head.weight": [153713, 2048] + } + } + } +} diff --git a/tools/hf2gguf/starvla/compare_starvla_actions.py b/tools/hf2gguf/starvla/compare_starvla_actions.py deleted file mode 100644 index e2422dc..0000000 --- a/tools/hf2gguf/starvla/compare_starvla_actions.py +++ /dev/null @@ -1,82 +0,0 @@ -#!/usr/bin/env python3 -"""Compare C++ StarVLA actions with a local Python reference.""" - -from __future__ import annotations - -import argparse -import json -import math -import sys -from pathlib import Path -from typing import Any - -import numpy as np - - -ACTION_RELATIVE_L2_LIMIT = 0.03 - - -class ComparisonError(RuntimeError): - pass - - -def load_actions(path: Path, key: str) -> np.ndarray: - try: - value: Any = json.loads(path.read_text(encoding="utf-8")) - for part in key.split("."): - value = value[part] - actions = np.asarray(value, dtype=np.float64) - except (OSError, json.JSONDecodeError, KeyError, TypeError, ValueError) as exc: - raise ComparisonError(f"cannot load {key!r} from {path}: {exc}") from exc - if actions.ndim == 3 and actions.shape[0] == 1: - actions = actions[0] - if actions.ndim != 2 or 0 in actions.shape: - raise ComparisonError(f"{path}:{key} must have shape [steps, dims]") - if not np.isfinite(actions).all(): - raise ComparisonError(f"{path}:{key} contains non-finite values") - return actions - - -def compare_actions(reference: np.ndarray, candidate: np.ndarray) -> dict[str, Any]: - if reference.shape != candidate.shape: - raise ComparisonError( - f"action shape mismatch: reference={reference.shape}, candidate={candidate.shape}" - ) - difference_l2 = float(np.linalg.norm(candidate - reference)) - reference_l2 = float(np.linalg.norm(reference)) - relative_l2 = difference_l2 / reference_l2 if reference_l2 else ( - 0.0 if difference_l2 == 0.0 else math.inf - ) - return { - "shape": list(reference.shape), - "relative_l2": relative_l2, - "limit": ACTION_RELATIVE_L2_LIMIT, - "passed": relative_l2 <= ACTION_RELATIVE_L2_LIMIT + 1e-12, - } - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--reference", type=Path, required=True) - parser.add_argument("--candidate", type=Path, required=True) - parser.add_argument("--reference-key", default="outputs.unnormalized_actions") - parser.add_argument("--candidate-key", default="actions") - return parser.parse_args() - - -def main() -> int: - args = parse_args() - try: - result = compare_actions( - load_actions(args.reference, args.reference_key), - load_actions(args.candidate, args.candidate_key), - ) - except ComparisonError as exc: - print(f"error: {exc}", file=sys.stderr) - return 2 - print(json.dumps(result, indent=2)) - return 0 if result["passed"] else 1 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/tools/hf2gguf/starvla/convert.sh b/tools/hf2gguf/starvla/convert.sh new file mode 100755 index 0000000..7606055 --- /dev/null +++ b/tools/hf2gguf/starvla/convert.sh @@ -0,0 +1,160 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../../.." && pwd)" +PYTHON="${PYTHON:-${ROOT_DIR}/.venv/bin/python}" +CATALOG="${ROOT_DIR}/tools/hf2gguf/starvla/checkpoint_catalog.json" + +usage() { + cat >&2 <<'EOF' +Usage: tools/hf2gguf/starvla/convert.sh VARIANT [OUTPUT_DIR] [OPTIONS] + +Options: + --checkpoint PATH Convert a training checkpoint instead of the catalog release + --source-dir DIR Run directory containing config.yaml and dataset_statistics.json + --unnorm-key KEY Default normalization profile for a training checkpoint +EOF +} + +if [[ $# -eq 1 && ( "$1" == -h || "$1" == --help ) ]]; then + usage + exit 0 +fi +if [[ $# -lt 1 ]]; then + usage + exit 2 +fi + +VARIANT=$1 +shift +OUTPUT_DIR="${ROOT_DIR}/ckpts/starvla/gguf/${VARIANT}" +if [[ $# -gt 0 && "$1" != --* ]]; then + OUTPUT_DIR=$1 + shift +fi +CHECKPOINT_OVERRIDE="" +SOURCE_DIR_OVERRIDE="" +UNNORM_KEY="" +while [[ $# -gt 0 ]]; do + case "$1" in + --checkpoint|--source-dir|--unnorm-key) + [[ $# -ge 2 ]] || { echo "error: $1 requires a value" >&2; exit 2; } + case "$1" in + --checkpoint) CHECKPOINT_OVERRIDE=$2 ;; + --source-dir) SOURCE_DIR_OVERRIDE=$2 ;; + --unnorm-key) UNNORM_KEY=$2 ;; + esac + shift 2 + ;; + -h|--help) + usage + exit 0 + ;; + *) + echo "error: unknown option: $1" >&2 + usage + exit 2 + ;; + esac +done +if [[ -z "${CHECKPOINT_OVERRIDE}" && ( -n "${SOURCE_DIR_OVERRIDE}" || -n "${UNNORM_KEY}" ) ]]; then + echo "error: --source-dir and --unnorm-key require --checkpoint" >&2 + exit 2 +fi +[[ -x "${PYTHON}" ]] || { echo "error: missing Python: ${PYTHON}" >&2; exit 2; } +[[ ! -e "${OUTPUT_DIR}" ]] || { + echo "error: refusing to overwrite output directory: ${OUTPUT_DIR}" >&2 + exit 2 +} + +export STARVLA_CONFIG_PYTHON="${PYTHON}" +source "${ROOT_DIR}/tools/hf2gguf/starvla/starvla_variant_config.sh" +load_starvla_variant "${VARIANT}" + +download_args=(--variant "${VARIANT}") +[[ "${FRAMEWORK}" == fast ]] && download_args+=(--include-fast-weights) +[[ -n "${CHECKPOINT_OVERRIDE}" ]] && download_args+=(--skip-checkpoint) +[[ "${STARVLA_LOCAL_FILES_ONLY:-0}" == 1 ]] && download_args+=(--local-files-only) +"${PYTHON}" "${ROOT_DIR}/tools/hf2gguf/starvla/download_starvla.py" \ + --catalog "${CATALOG}" "${download_args[@]}" + +LLAMA_REV="$("${PYTHON}" -c \ + 'import json,sys; print(json.load(open(sys.argv[1]))["source_revisions"]["llama_cpp"])' \ + "${CATALOG}")" +export LLAMA_ROOT="${LLAMA_ROOT:-${ROOT_DIR}/ckpts/starvla/toolchains/llama.cpp-${LLAMA_REV}}" +if [[ ! -d "${LLAMA_ROOT}" ]]; then + mkdir -p -- "$(dirname -- "${LLAMA_ROOT}")" + git -C "${ROOT_DIR}/third_party/llama.cpp" worktree add \ + --detach "${LLAMA_ROOT}" "${LLAMA_REV}" +fi + +SOURCE_DIR="${ROOT_DIR}/ckpts/starvla/sources/${CHECKPOINT_DIRECTORY}/${CHECKPOINT_REVISION}" +CHECKPOINT="${SOURCE_DIR}/${CHECKPOINT_RELATIVE_PATH}" +BASE_ASSETS="${ROOT_DIR}/ckpts/starvla/sources/${QWEN_DIRECTORY}/${QWEN_REVISION}" +mkdir -p -- "${ROOT_DIR}/ckpts/starvla/work" +WORK_DIR="$(mktemp -d "${ROOT_DIR}/ckpts/starvla/work/.${VARIANT}.XXXXXX")" +cleanup() { rm -rf -- "${WORK_DIR}"; } +trap cleanup EXIT +trap 'exit 130' INT +trap 'exit 143' TERM + +if [[ -n "${CHECKPOINT_OVERRIDE}" ]]; then + CHECKPOINT="$(realpath -- "${CHECKPOINT_OVERRIDE}")" + if [[ -n "${SOURCE_DIR_OVERRIDE}" ]]; then + SOURCE_DIR="$(realpath -- "${SOURCE_DIR_OVERRIDE}")" + else + checkpoint_dir="$(dirname -- "${CHECKPOINT}")" + for candidate in "${checkpoint_dir}" "$(dirname -- "${checkpoint_dir}")"; do + if [[ -f "${candidate}/config.yaml" && -f "${candidate}/dataset_statistics.json" ]]; then + SOURCE_DIR="${candidate}" + break + fi + done + fi + [[ -f "${SOURCE_DIR}/config.yaml" && -f "${SOURCE_DIR}/dataset_statistics.json" ]] || { + echo "error: cannot find config.yaml and dataset_statistics.json; pass --source-dir" >&2 + exit 2 + } + LOCAL_CATALOG="${WORK_DIR}/checkpoint_catalog.json" + "${PYTHON}" - "${CATALOG}" "${LOCAL_CATALOG}" "${VARIANT}" \ + "${CHECKPOINT}" "${SOURCE_DIR}" "${UNNORM_KEY}" <<'PY' +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(sys.argv[1]).parent)) +from starvla_checkpoint import atomic_write_json, load_catalog, local_checkpoint_catalog + +catalog = local_checkpoint_catalog( + load_catalog(Path(sys.argv[1])), + sys.argv[3], + Path(sys.argv[4]), + Path(sys.argv[5]), + sys.argv[6] or None, +) +atomic_write_json(Path(sys.argv[2]), catalog) +PY + CATALOG="${LOCAL_CATALOG}" + export STARVLA_CATALOG="${CATALOG}" + load_starvla_variant "${VARIANT}" +fi + +export PYTHON VARIANT SOURCE_DIR CHECKPOINT BASE_ASSETS WORK_DIR OUTPUT_DIR LLAMA_ROOT +if [[ "${FRAMEWORK}" != fast ]]; then + bash "${ROOT_DIR}/tools/hf2gguf/starvla/convert_starvla_all.sh" +else + CODEC_REV="$("${PYTHON}" -c \ + 'import json,sys; print(json.load(open(sys.argv[1]))["shared_assets"]["fast_codec"]["revision"])' \ + "${CATALOG}")" + "${PYTHON}" "${ROOT_DIR}/tools/hf2gguf/starvla/convert_starvla_qwen25_fast.py" \ + --checkpoint "${CHECKPOINT}" \ + --source-dir "${SOURCE_DIR}" \ + --qwen-assets "${BASE_ASSETS}" \ + --fast-codec "${ROOT_DIR}/ckpts/starvla/sources/fast-codec/${CODEC_REV}" \ + --staging-dir "${WORK_DIR}/staging" \ + --output-dir "${OUTPUT_DIR}" \ + --catalog "${CATALOG}" \ + --llama-root "${LLAMA_ROOT}" \ + --python "${PYTHON}" +fi + +echo "StarVLA ${VARIANT} bundle: ${OUTPUT_DIR}" diff --git a/tools/hf2gguf/starvla/convert_starvla_all.sh b/tools/hf2gguf/starvla/convert_starvla_all.sh new file mode 100755 index 0000000..a0fa6c5 --- /dev/null +++ b/tools/hf2gguf/starvla/convert_starvla_all.sh @@ -0,0 +1,144 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)" +PYTHON="${PYTHON:-${ROOT_DIR}/.venv/bin/python}" +CATALOG="${STARVLA_CATALOG:-${ROOT_DIR}/tools/hf2gguf/starvla/checkpoint_catalog.json}" +LLAMA_ROOT="${LLAMA_ROOT:?set LLAMA_ROOT to an absolute clean checkout of the catalog-pinned llama.cpp revision}" +source "${ROOT_DIR}/tools/hf2gguf/starvla/starvla_variant_config.sh" +VARIANT="${VARIANT:?set VARIANT to oft, groot, pi_v3, qwen25_oft, qwen25_groot, or qwen25_pi}" +load_starvla_variant "${VARIANT}" + +if [[ "${FRAMEWORK}" == "fast" ]]; then + echo "error: FAST uses tools/hf2gguf/starvla/convert_starvla_qwen25_fast.py" >&2 + exit 1 +fi + +CHECKPOINT="${CHECKPOINT:?set CHECKPOINT to the pinned ${VARIANT} .pt file}" +SOURCE_DIR="${SOURCE_DIR:?set SOURCE_DIR to the pinned ${VARIANT} source directory}" +BASE_ASSETS="${BASE_ASSETS:?set BASE_ASSETS to the pinned Qwen-VL asset directory}" +WORK_DIR="${WORK_DIR:-${ROOT_DIR}/ckpts/starvla/work/${VARIANT}}" +OUTPUT_DIR="${OUTPUT_DIR:-${ROOT_DIR}/ckpts/starvla/gguf/${VARIANT}}" +MAX_SHARD_SIZE="${MAX_SHARD_SIZE:-2G}" +TEXT_DTYPE="${TEXT_DTYPE:-bf16}" +MMPROJ_DTYPE="${MMPROJ_DTYPE:-bf16}" +POLICY_DTYPE="${POLICY_DTYPE:-fp32}" +TEXT_FILENAME="${TEXT_FILENAME:-qwen-${ARTIFACT_STEM}-${TEXT_DTYPE}.gguf}" +MMPROJ_FILENAME="${MMPROJ_FILENAME:-mmproj-${ARTIFACT_STEM}-${MMPROJ_DTYPE}.gguf}" +POLICY_FILENAME="${POLICY_FILENAME:-starvla-${ARTIFACT_STEM}-policy-${POLICY_DTYPE}.gguf}" +MANIFEST_FILENAME="conversion_manifest.json" + +if [[ -e "${WORK_DIR}" ]] && + [[ -n "$(find "${WORK_DIR}" -mindepth 1 -maxdepth 1 -print -quit)" ]]; then + echo "error: WORK_DIR must be empty: ${WORK_DIR}" >&2 + exit 1 +fi + +mkdir -p "${OUTPUT_DIR}" + +for filename in \ + "${TEXT_FILENAME}" \ + "${MMPROJ_FILENAME}" \ + "${POLICY_FILENAME}" \ + "${MANIFEST_FILENAME}"; do + destination="${OUTPUT_DIR}/${filename}" + if [[ -e "${destination}" || -L "${destination}" ]]; then + echo "error: refusing to overwrite existing output: ${destination}" >&2 + exit 1 + fi +done + +RUN_OUTPUT_DIR="" +declare -a PUBLISHED_FILES=() +SUCCESS=0 + +cleanup() { + status=$? + trap - EXIT + set +e + if [[ -n "${RUN_OUTPUT_DIR}" ]]; then + rm -rf -- "${RUN_OUTPUT_DIR}" + fi + if [[ "${SUCCESS}" != 1 ]]; then + for published in "${PUBLISHED_FILES[@]}"; do + rm -f -- "${published}" + done + rm -rf -- "${WORK_DIR}" + fi + exit "${status}" +} +trap cleanup EXIT +trap 'exit 130' INT +trap 'exit 143' TERM + +mkdir -p "${WORK_DIR}" +RUN_OUTPUT_DIR="$(mktemp -d "${OUTPUT_DIR}/.starvla-${ARTIFACT_STEM}.tmp.XXXXXX")" + +publish_file() { + source=$1 + destination=$2 + if [[ ! -f "${source}" || ! -s "${source}" ]]; then + echo "error: transaction output is missing or empty: ${source}" >&2 + return 1 + fi + if ! ln -- "${source}" "${destination}"; then + echo "error: refusing to overwrite existing output: ${destination}" >&2 + return 1 + fi + PUBLISHED_FILES+=("${destination}") + rm -f -- "${source}" +} + +"${PYTHON}" "${ROOT_DIR}/tools/hf2gguf/starvla/starvla_surgery.py" \ + "${CHECKPOINT}" \ + --variant "${VARIANT}" \ + --catalog "${CATALOG}" \ + --source-dir "${SOURCE_DIR}" \ + --base-assets "${BASE_ASSETS}" \ + --output-dir "${WORK_DIR}/staging" \ + --max-shard-size "${MAX_SHARD_SIZE}" + +"${PYTHON}" "${ROOT_DIR}/tools/hf2gguf/starvla/convert_starvla_qwen_to_gguf.py" \ + --hf-dir "${WORK_DIR}/staging/hf" \ + --surgery-manifest "${WORK_DIR}/staging/surgery_manifest.json" \ + --catalog "${CATALOG}" \ + --output-dir "${RUN_OUTPUT_DIR}" \ + --llama-root "${LLAMA_ROOT}" \ + --text-filename "${TEXT_FILENAME}" \ + --mmproj-filename "${MMPROJ_FILENAME}" \ + --text-dtype "${TEXT_DTYPE}" \ + --mmproj-dtype "${MMPROJ_DTYPE}" + +"${PYTHON}" "${ROOT_DIR}/tools/hf2gguf/starvla/convert_starvla_policy_to_gguf.py" \ + --variant "${VARIANT}" \ + --policy-dir "${WORK_DIR}/staging/policy" \ + --hf-dir "${WORK_DIR}/staging/hf" \ + --surgery-manifest "${WORK_DIR}/staging/surgery_manifest.json" \ + --catalog "${CATALOG}" \ + --output "${RUN_OUTPUT_DIR}/${POLICY_FILENAME}" \ + --dtype "${POLICY_DTYPE}" \ + --text-filename "${TEXT_FILENAME}" \ + --mmproj-filename "${MMPROJ_FILENAME}" + +"${PYTHON}" "${ROOT_DIR}/tools/hf2gguf/starvla/validate_starvla_bundle.py" \ + --variant "${VARIANT}" \ + --text "${RUN_OUTPUT_DIR}/${TEXT_FILENAME}" \ + --mmproj "${RUN_OUTPUT_DIR}/${MMPROJ_FILENAME}" \ + --policy "${RUN_OUTPUT_DIR}/${POLICY_FILENAME}" \ + --hf-dir "${WORK_DIR}/staging/hf" \ + --policy-dir "${WORK_DIR}/staging/policy" \ + --surgery-manifest "${WORK_DIR}/staging/surgery_manifest.json" \ + --catalog "${CATALOG}" \ + --text-dtype "${TEXT_DTYPE}" \ + --mmproj-dtype "${MMPROJ_DTYPE}" \ + --policy-dtype "${POLICY_DTYPE}" \ + --output "${RUN_OUTPUT_DIR}/${MANIFEST_FILENAME}" + +publish_file "${RUN_OUTPUT_DIR}/${TEXT_FILENAME}" "${OUTPUT_DIR}/${TEXT_FILENAME}" +publish_file "${RUN_OUTPUT_DIR}/${MMPROJ_FILENAME}" "${OUTPUT_DIR}/${MMPROJ_FILENAME}" +publish_file "${RUN_OUTPUT_DIR}/${POLICY_FILENAME}" "${OUTPUT_DIR}/${POLICY_FILENAME}" +# The manifest is the bundle commit marker and is intentionally published last. +publish_file "${RUN_OUTPUT_DIR}/${MANIFEST_FILENAME}" "${OUTPUT_DIR}/${MANIFEST_FILENAME}" +SUCCESS=1 + +echo "StarVLA ${VARIANT} bundle written to ${OUTPUT_DIR}" diff --git a/tools/hf2gguf/starvla/convert_starvla_policy_to_gguf.py b/tools/hf2gguf/starvla/convert_starvla_policy_to_gguf.py index c8c53ed..21dd35a 100755 --- a/tools/hf2gguf/starvla/convert_starvla_policy_to_gguf.py +++ b/tools/hf2gguf/starvla/convert_starvla_policy_to_gguf.py @@ -4,14 +4,12 @@ from __future__ import annotations import argparse -import copy -import hashlib import json import math import os import sys from pathlib import Path -from typing import Any +from typing import Any, Mapping import numpy as np @@ -28,9 +26,8 @@ load_catalog, resolve_effective_config, sha256_file, - validate_official_surgery_manifest, + validate_surgery_manifest, verify_staged_assets, - verify_staged_tensors_against_checkpoint, ) @@ -54,117 +51,77 @@ } -def build_groot_tensor_map(block_count: int = 16) -> dict[str, str]: - """Return the complete released Qwen-GR00T policy tensor renaming map.""" +DIT_BLOCK_SUFFIXES = { + "norm1.linear.weight": "ada_norm.weight", + "norm1.linear.bias": "ada_norm.bias", + "attn1.to_q.weight": "attention.query.weight", + "attn1.to_q.bias": "attention.query.bias", + "attn1.to_k.weight": "attention.key.weight", + "attn1.to_k.bias": "attention.key.bias", + "attn1.to_v.weight": "attention.value.weight", + "attn1.to_v.bias": "attention.value.bias", + "attn1.to_out.0.weight": "attention.output.weight", + "attn1.to_out.0.bias": "attention.output.bias", + "ff.net.0.proj.weight": "feed_forward.input.weight", + "ff.net.0.proj.bias": "feed_forward.input.bias", + "ff.net.2.weight": "feed_forward.output.weight", + "ff.net.2.bias": "feed_forward.output.bias", +} + + +def _build_flow_tensor_map(framework: str, block_count: int) -> dict[str, str]: + destination = f"starvla.policy.{framework}" tensor_map = { "action_model.model.timestep_encoder.timestep_embedder.linear_1.weight": - "starvla.policy.groot.timestep.input.weight", + f"{destination}.timestep.input.weight", "action_model.model.timestep_encoder.timestep_embedder.linear_1.bias": - "starvla.policy.groot.timestep.input.bias", + f"{destination}.timestep.input.bias", "action_model.model.timestep_encoder.timestep_embedder.linear_2.weight": - "starvla.policy.groot.timestep.output.weight", + f"{destination}.timestep.output.weight", "action_model.model.timestep_encoder.timestep_embedder.linear_2.bias": - "starvla.policy.groot.timestep.output.bias", - } - block_suffixes = { - "norm1.linear.weight": "ada_norm.weight", - "norm1.linear.bias": "ada_norm.bias", - "attn1.to_q.weight": "attention.query.weight", - "attn1.to_q.bias": "attention.query.bias", - "attn1.to_k.weight": "attention.key.weight", - "attn1.to_k.bias": "attention.key.bias", - "attn1.to_v.weight": "attention.value.weight", - "attn1.to_v.bias": "attention.value.bias", - "attn1.to_out.0.weight": "attention.output.weight", - "attn1.to_out.0.bias": "attention.output.bias", - "ff.net.0.proj.weight": "feed_forward.input.weight", - "ff.net.0.proj.bias": "feed_forward.input.bias", - "ff.net.2.weight": "feed_forward.output.weight", - "ff.net.2.bias": "feed_forward.output.bias", + f"{destination}.timestep.output.bias", + "action_model.action_encoder.layer1.weight": f"{destination}.action.input.weight", + "action_model.action_encoder.layer1.bias": f"{destination}.action.input.bias", + "action_model.action_encoder.layer2.weight": f"{destination}.action.time_mix.weight", + "action_model.action_encoder.layer2.bias": f"{destination}.action.time_mix.bias", + "action_model.action_encoder.layer3.weight": f"{destination}.action.output.weight", + "action_model.action_encoder.layer3.bias": f"{destination}.action.output.bias", + "action_model.action_decoder.layer1.weight": f"{destination}.velocity.input.weight", + "action_model.action_decoder.layer1.bias": f"{destination}.velocity.input.bias", + "action_model.action_decoder.layer2.weight": f"{destination}.velocity.output.weight", + "action_model.action_decoder.layer2.bias": f"{destination}.velocity.output.bias", + "action_model.future_tokens.weight": f"{destination}.future_tokens.weight", + "action_model.position_embedding.weight": f"{destination}.action_position.weight", } for block in range(block_count): - for source_suffix, destination_suffix in block_suffixes.items(): - tensor_map[f"action_model.model.transformer_blocks.{block}.{source_suffix}"] = ( - f"starvla.policy.groot.block.{block}.{destination_suffix}" - ) + source = f"action_model.model.transformer_blocks.{block}" + target = f"{destination}.block.{block}" + for source_suffix, destination_suffix in DIT_BLOCK_SUFFIXES.items(): + tensor_map[f"{source}.{source_suffix}"] = f"{target}.{destination_suffix}" + return tensor_map + + +def build_groot_tensor_map(block_count: int = 16) -> dict[str, str]: + tensor_map = _build_flow_tensor_map("groot", block_count) tensor_map.update( { "action_model.model.proj_out_1.weight": "starvla.policy.groot.output.modulation.weight", "action_model.model.proj_out_1.bias": "starvla.policy.groot.output.modulation.bias", "action_model.model.proj_out_2.weight": "starvla.policy.groot.output.projection.weight", "action_model.model.proj_out_2.bias": "starvla.policy.groot.output.projection.bias", - "action_model.action_encoder.layer1.weight": "starvla.policy.groot.action.input.weight", - "action_model.action_encoder.layer1.bias": "starvla.policy.groot.action.input.bias", - "action_model.action_encoder.layer2.weight": "starvla.policy.groot.action.time_mix.weight", - "action_model.action_encoder.layer2.bias": "starvla.policy.groot.action.time_mix.bias", - "action_model.action_encoder.layer3.weight": "starvla.policy.groot.action.output.weight", - "action_model.action_encoder.layer3.bias": "starvla.policy.groot.action.output.bias", - "action_model.action_decoder.layer1.weight": "starvla.policy.groot.velocity.input.weight", - "action_model.action_decoder.layer1.bias": "starvla.policy.groot.velocity.input.bias", - "action_model.action_decoder.layer2.weight": "starvla.policy.groot.velocity.output.weight", - "action_model.action_decoder.layer2.bias": "starvla.policy.groot.velocity.output.bias", - "action_model.future_tokens.weight": "starvla.policy.groot.future_tokens.weight", - "action_model.position_embedding.weight": "starvla.policy.groot.action_position.weight", } ) return tensor_map def build_pi_tensor_map(block_count: int = 16) -> dict[str, str]: - """Return tensors used by the legacy Qwen-PI inference graph.""" - tensor_map = { - "action_model.model.timestep_encoder.timestep_embedder.linear_1.weight": - "starvla.policy.pi.timestep.input.weight", - "action_model.model.timestep_encoder.timestep_embedder.linear_1.bias": - "starvla.policy.pi.timestep.input.bias", - "action_model.model.timestep_encoder.timestep_embedder.linear_2.weight": - "starvla.policy.pi.timestep.output.weight", - "action_model.model.timestep_encoder.timestep_embedder.linear_2.bias": - "starvla.policy.pi.timestep.output.bias", - } - block_suffixes = { - "norm1.linear.weight": "ada_norm.weight", - "norm1.linear.bias": "ada_norm.bias", - "attn1.to_q.weight": "attention.query.weight", - "attn1.to_q.bias": "attention.query.bias", - "attn1.to_k.weight": "attention.key.weight", - "attn1.to_k.bias": "attention.key.bias", - "attn1.to_v.weight": "attention.value.weight", - "attn1.to_v.bias": "attention.value.bias", - "attn1.to_out.0.weight": "attention.output.weight", - "attn1.to_out.0.bias": "attention.output.bias", - "ff.net.0.proj.weight": "feed_forward.input.weight", - "ff.net.0.proj.bias": "feed_forward.input.bias", - "ff.net.2.weight": "feed_forward.output.weight", - "ff.net.2.bias": "feed_forward.output.bias", - } - for block in range(block_count): - for source_suffix, destination_suffix in block_suffixes.items(): - tensor_map[f"action_model.model.transformer_blocks.{block}.{source_suffix}"] = ( - f"starvla.policy.pi.block.{block}.{destination_suffix}" - ) + tensor_map = _build_flow_tensor_map("pi", block_count) tensor_map.update( { "action_model.state_encoder.layer1.weight": "starvla.policy.pi.state.input.weight", "action_model.state_encoder.layer1.bias": "starvla.policy.pi.state.input.bias", "action_model.state_encoder.layer2.weight": "starvla.policy.pi.state.output.weight", "action_model.state_encoder.layer2.bias": "starvla.policy.pi.state.output.bias", - "action_model.action_encoder.layer1.weight": "starvla.policy.pi.action.input.weight", - "action_model.action_encoder.layer1.bias": "starvla.policy.pi.action.input.bias", - "action_model.action_encoder.layer2.weight": - "starvla.policy.pi.action.time_mix.weight", - "action_model.action_encoder.layer2.bias": - "starvla.policy.pi.action.time_mix.bias", - "action_model.action_encoder.layer3.weight": "starvla.policy.pi.action.output.weight", - "action_model.action_encoder.layer3.bias": "starvla.policy.pi.action.output.bias", - "action_model.action_decoder.layer1.weight": - "starvla.policy.pi.velocity.input.weight", - "action_model.action_decoder.layer1.bias": "starvla.policy.pi.velocity.input.bias", - "action_model.action_decoder.layer2.weight": - "starvla.policy.pi.velocity.output.weight", - "action_model.action_decoder.layer2.bias": "starvla.policy.pi.velocity.output.bias", - "action_model.future_tokens.weight": "starvla.policy.pi.future_tokens.weight", - "action_model.position_embedding.weight": "starvla.policy.pi.action_position.weight", } ) return tensor_map @@ -174,63 +131,16 @@ def build_pi_v3_tensor_map( block_count: int = 36, projector_count: int = 36, ) -> dict[str, str]: - """Return the tensors used by the Qwen PI-v3 inference graph.""" - tensor_map = { - "action_model.model.timestep_encoder.timestep_embedder.linear_1.weight": - "starvla.policy.pi_v3.timestep.input.weight", - "action_model.model.timestep_encoder.timestep_embedder.linear_1.bias": - "starvla.policy.pi_v3.timestep.input.bias", - "action_model.model.timestep_encoder.timestep_embedder.linear_2.weight": - "starvla.policy.pi_v3.timestep.output.weight", - "action_model.model.timestep_encoder.timestep_embedder.linear_2.bias": - "starvla.policy.pi_v3.timestep.output.bias", - } - block_suffixes = { - "norm1.linear.weight": "ada_norm.weight", - "norm1.linear.bias": "ada_norm.bias", - "attn1.to_q.weight": "attention.query.weight", - "attn1.to_q.bias": "attention.query.bias", - "attn1.to_k.weight": "attention.key.weight", - "attn1.to_k.bias": "attention.key.bias", - "attn1.to_v.weight": "attention.value.weight", - "attn1.to_v.bias": "attention.value.bias", - "attn1.to_out.0.weight": "attention.output.weight", - "attn1.to_out.0.bias": "attention.output.bias", - "ff.net.0.proj.weight": "feed_forward.input.weight", - "ff.net.0.proj.bias": "feed_forward.input.bias", - "ff.net.2.weight": "feed_forward.output.weight", - "ff.net.2.bias": "feed_forward.output.bias", - } - for block in range(block_count): - for source_suffix, destination_suffix in block_suffixes.items(): - tensor_map[f"action_model.model.transformer_blocks.{block}.{source_suffix}"] = ( - f"starvla.policy.pi_v3.block.{block}.{destination_suffix}" - ) - tensor_map.update( - { - "action_model.action_encoder.layer1.weight": "starvla.policy.pi_v3.action.input.weight", - "action_model.action_encoder.layer1.bias": "starvla.policy.pi_v3.action.input.bias", - "action_model.action_encoder.layer2.weight": "starvla.policy.pi_v3.action.time_mix.weight", - "action_model.action_encoder.layer2.bias": "starvla.policy.pi_v3.action.time_mix.bias", - "action_model.action_encoder.layer3.weight": "starvla.policy.pi_v3.action.output.weight", - "action_model.action_encoder.layer3.bias": "starvla.policy.pi_v3.action.output.bias", - "action_model.action_decoder.layer1.weight": "starvla.policy.pi_v3.velocity.input.weight", - "action_model.action_decoder.layer1.bias": "starvla.policy.pi_v3.velocity.input.bias", - "action_model.action_decoder.layer2.weight": "starvla.policy.pi_v3.velocity.output.weight", - "action_model.action_decoder.layer2.bias": "starvla.policy.pi_v3.velocity.output.bias", - "action_model.future_tokens.weight": "starvla.policy.pi_v3.future_tokens.weight", - "action_model.position_embedding.weight": "starvla.policy.pi_v3.action_position.weight", - } - ) + tensor_map = _build_flow_tensor_map("pi_v3", block_count) for projector in range(projector_count): - source_prefix = f"project_layers.{projector}" - destination_prefix = f"starvla.policy.pi_v3.projector.{projector}" + source = f"project_layers.{projector}" + target = f"starvla.policy.pi_v3.projector.{projector}" tensor_map.update( { - f"{source_prefix}.0.weight": f"{destination_prefix}.norm.weight", - f"{source_prefix}.0.bias": f"{destination_prefix}.norm.bias", - f"{source_prefix}.1.weight": f"{destination_prefix}.projection.weight", - f"{source_prefix}.1.bias": f"{destination_prefix}.projection.bias", + f"{source}.0.weight": f"{target}.norm.weight", + f"{source}.0.bias": f"{target}.norm.bias", + f"{source}.1.weight": f"{target}.projection.weight", + f"{source}.1.bias": f"{target}.projection.bias", } ) return tensor_map @@ -250,7 +160,7 @@ def build_pi_v3_tensor_map( GROOT_QWEN25_POLICY_NUMEL = 155_181_319 GROOT_DIT_NORM_EPS = 1e-5 GROOT_OUTPUT_NORM_EPS = 1e-6 -GROOT_OFFICIAL_DIMENSIONS_BY_BACKBONE = { +GROOT_SUPPORTED_DIMENSIONS_BY_BACKBONE = { backbone: { "qwen_hidden_dim": qwen_hidden_dim, "dit_width": 768, @@ -271,7 +181,6 @@ def build_pi_v3_tensor_map( ("qwen2_5_vl", 2048, GROOT_QWEN25_POLICY_NUMEL), ) } -GROOT_OFFICIAL_DIMENSIONS = GROOT_OFFICIAL_DIMENSIONS_BY_BACKBONE["qwen3_vl"] PI_BLOCK_COUNT = 16 PI_TENSOR_MAP = build_pi_tensor_map(PI_BLOCK_COUNT) @@ -284,7 +193,7 @@ def build_pi_v3_tensor_map( PI_POLICY_TENSOR_COUNT = 244 PI_POLICY_NUMEL = 967_796_743 PI_DIT_NORM_EPS = 1e-5 -PI_OFFICIAL_DIMENSIONS = { +PI_SUPPORTED_DIMENSIONS = { "qwen_hidden_dim": 2048, "dit_width": 2048, "timestep_dim": 256, @@ -305,7 +214,7 @@ def build_pi_v3_tensor_map( PI_V3_POLICY_TENSOR_COUNT = len(PI_V3_TENSOR_MAP) PI_V3_DIT_NORM_EPS = 1e-5 PI_V3_PROJECTOR_NORM_EPS = 1e-5 -PI_V3_OFFICIAL_DIMENSIONS = { +PI_V3_SUPPORTED_DIMENSIONS = { "qwen_hidden_dim": 2560, "dit_width": 1024, "timestep_dim": 256, @@ -453,6 +362,50 @@ def load_policy_tensors(policy_dir: Path) -> dict[str, Any]: return tensors +def _tensor_shape(tensors: Mapping[str, Any], name: str) -> list[int]: + return [int(dimension) for dimension in tensors[name].shape] + + +def _matrix_shape(tensors: Mapping[str, Any], name: str) -> list[int]: + shape = _tensor_shape(tensors, name) + if len(shape) != 2: + raise StarVLAError(f"invalid matrix shape for {name}: {shape}") + return shape + + +def _validate_tensor_shapes( + tensors: Mapping[str, Any], expected: Mapping[str, list[int]], *, label: str +) -> None: + mismatches = [ + f"{name}: expected {shape}, got {_tensor_shape(tensors, name)}" + for name, shape in expected.items() + if _tensor_shape(tensors, name) != shape + ] + if mismatches: + raise StarVLAError(f"invalid {label} tensor shapes: " + "; ".join(mismatches)) + + +def _dit_block_shapes( + prefix: str, width: int, attention_dim: int, feed_forward_dim: int +) -> dict[str, list[int]]: + return { + f"{prefix}.norm1.linear.weight": [2 * width, width], + f"{prefix}.norm1.linear.bias": [2 * width], + f"{prefix}.attn1.to_q.weight": [width, width], + f"{prefix}.attn1.to_q.bias": [width], + f"{prefix}.attn1.to_k.weight": [width, attention_dim], + f"{prefix}.attn1.to_k.bias": [width], + f"{prefix}.attn1.to_v.weight": [width, attention_dim], + f"{prefix}.attn1.to_v.bias": [width], + f"{prefix}.attn1.to_out.0.weight": [width, width], + f"{prefix}.attn1.to_out.0.bias": [width], + f"{prefix}.ff.net.0.proj.weight": [feed_forward_dim, width], + f"{prefix}.ff.net.0.proj.bias": [feed_forward_dim], + f"{prefix}.ff.net.2.weight": [width, feed_forward_dim], + f"{prefix}.ff.net.2.bias": [width], + } + + def validate_oft_tensors(tensors: dict[str, Any]) -> dict[str, int]: actual = set(tensors) expected = set(OFT_TENSOR_MAP) @@ -461,15 +414,12 @@ def validate_oft_tensors(tensors: dict[str, Any]) -> dict[str, int]: unexpected = sorted(actual - expected) raise StarVLAError(f"OFT policy tensor mismatch; missing={missing}, unexpected={unexpected}") - def shape(name: str) -> list[int]: - return [int(dim) for dim in tensors[name].shape] - - input_dim = shape("action_model.model.layer_norm1.weight")[0] - input_projection = shape("action_model.model.fc1.weight") + input_dim = _tensor_shape(tensors, "action_model.model.layer_norm1.weight")[0] + input_projection = _tensor_shape(tensors, "action_model.model.fc1.weight") if len(input_projection) != 2 or input_projection[1] != input_dim: raise StarVLAError(f"invalid OFT input projection shape: {input_projection}") hidden_dim = input_projection[0] - output_projection = shape("action_model.model.fc2.weight") + output_projection = _tensor_shape(tensors, "action_model.model.fc2.weight") if len(output_projection) != 2 or output_projection[1] != hidden_dim: raise StarVLAError(f"invalid OFT output projection shape: {output_projection}") action_dim = output_projection[0] @@ -492,18 +442,12 @@ def shape(name: str) -> list[int]: "action_model.model.fc2.weight": [action_dim, hidden_dim], "action_model.model.fc2.bias": [action_dim], } - mismatches = [ - f"{name}: expected {expected_shape}, got {shape(name)}" - for name, expected_shape in expected_shapes.items() - if shape(name) != expected_shape - ] - if mismatches: - raise StarVLAError("invalid OFT tensor shapes: " + "; ".join(mismatches)) + _validate_tensor_shapes(tensors, expected_shapes, label="OFT") return {"input_dim": input_dim, "hidden_dim": hidden_dim, "action_dim": action_dim} def validate_groot_tensors(tensors: dict[str, Any]) -> dict[str, int]: - """Validate every released GR00T policy tensor and infer its architecture.""" + """Validate GR00T policy tensors and infer their dimensions.""" actual = set(tensors) expected = set(GROOT_TENSOR_MAP) if not expected.issubset(actual) or actual - expected != GROOT_UNUSED_SOURCE_TENSORS: @@ -511,31 +455,22 @@ def validate_groot_tensors(tensors: dict[str, Any]) -> dict[str, int]: unexpected = sorted(actual - expected - GROOT_UNUSED_SOURCE_TENSORS) raise StarVLAError(f"GR00T policy tensor mismatch; missing={missing}, unexpected={unexpected}") - def shape(name: str) -> list[int]: - return [int(dim) for dim in tensors[name].shape] - - def matrix_shape(name: str) -> list[int]: - value = shape(name) - if len(value) != 2: - raise StarVLAError(f"invalid GR00T matrix shape for {name}: {value}") - return value - - timestep_input = matrix_shape( + timestep_input = _matrix_shape(tensors, "action_model.model.timestep_encoder.timestep_embedder.linear_1.weight" ) dit_width, timestep_dim = timestep_input - cross_attention_dim = matrix_shape( + cross_attention_dim = _matrix_shape(tensors, "action_model.model.transformer_blocks.0.attn1.to_k.weight" )[1] - feed_forward_dim = matrix_shape( + feed_forward_dim = _matrix_shape(tensors, "action_model.model.transformer_blocks.0.ff.net.0.proj.weight" )[0] - output_dim = matrix_shape("action_model.model.proj_out_2.weight")[0] - mlp_hidden_dim = matrix_shape("action_model.action_decoder.layer1.weight")[0] - state_dim = matrix_shape("action_model.state_encoder.layer1.weight")[1] - action_dim = matrix_shape("action_model.action_encoder.layer1.weight")[1] - future_token_count = matrix_shape("action_model.future_tokens.weight")[0] - max_sequence_length = matrix_shape("action_model.position_embedding.weight")[0] + output_dim = _matrix_shape(tensors, "action_model.model.proj_out_2.weight")[0] + mlp_hidden_dim = _matrix_shape(tensors, "action_model.action_decoder.layer1.weight")[0] + state_dim = _matrix_shape(tensors, "action_model.state_encoder.layer1.weight")[1] + action_dim = _matrix_shape(tensors, "action_model.action_encoder.layer1.weight")[1] + future_token_count = _matrix_shape(tensors, "action_model.future_tokens.weight")[0] + max_sequence_length = _matrix_shape(tensors, "action_model.position_embedding.weight")[0] expected_shapes = { "action_model.model.timestep_encoder.timestep_embedder.linear_1.weight": [ @@ -566,33 +501,16 @@ def matrix_shape(name: str) -> list[int]: "action_model.position_embedding.weight": [max_sequence_length, dit_width], } for block in range(GROOT_BLOCK_COUNT): - prefix = f"action_model.model.transformer_blocks.{block}" - attention_input_dim = cross_attention_dim if block % 2 == 0 else dit_width + attention_dim = cross_attention_dim if block % 2 == 0 else dit_width expected_shapes.update( - { - f"{prefix}.norm1.linear.weight": [2 * dit_width, dit_width], - f"{prefix}.norm1.linear.bias": [2 * dit_width], - f"{prefix}.attn1.to_q.weight": [dit_width, dit_width], - f"{prefix}.attn1.to_q.bias": [dit_width], - f"{prefix}.attn1.to_k.weight": [dit_width, attention_input_dim], - f"{prefix}.attn1.to_k.bias": [dit_width], - f"{prefix}.attn1.to_v.weight": [dit_width, attention_input_dim], - f"{prefix}.attn1.to_v.bias": [dit_width], - f"{prefix}.attn1.to_out.0.weight": [dit_width, dit_width], - f"{prefix}.attn1.to_out.0.bias": [dit_width], - f"{prefix}.ff.net.0.proj.weight": [feed_forward_dim, dit_width], - f"{prefix}.ff.net.0.proj.bias": [feed_forward_dim], - f"{prefix}.ff.net.2.weight": [dit_width, feed_forward_dim], - f"{prefix}.ff.net.2.bias": [dit_width], - } + _dit_block_shapes( + f"action_model.model.transformer_blocks.{block}", + dit_width, + attention_dim, + feed_forward_dim, + ) ) - mismatches = [ - f"{name}: expected {expected_shape}, got {shape(name)}" - for name, expected_shape in expected_shapes.items() - if shape(name) != expected_shape - ] - if mismatches: - raise StarVLAError("invalid GR00T tensor shapes: " + "; ".join(mismatches)) + _validate_tensor_shapes(tensors, expected_shapes, label="GR00T") numel = sum(int(tensor.numel()) for tensor in tensors.values()) return { @@ -623,29 +541,20 @@ def validate_pi_tensors(tensors: dict[str, Any]) -> dict[str, int]: f"legacy PI policy tensor mismatch; missing={missing}, unexpected={unexpected}" ) - def shape(name: str) -> list[int]: - return [int(dim) for dim in tensors[name].shape] - - def matrix_shape(name: str) -> list[int]: - value = shape(name) - if len(value) != 2: - raise StarVLAError(f"invalid legacy PI matrix shape for {name}: {value}") - return value - - timestep_input = matrix_shape( + timestep_input = _matrix_shape(tensors, "action_model.model.timestep_encoder.timestep_embedder.linear_1.weight" ) dit_width, timestep_dim = timestep_input - cross_attention_dim = matrix_shape( + cross_attention_dim = _matrix_shape(tensors, "action_model.model.transformer_blocks.0.attn1.to_k.weight" )[1] - feed_forward_dim = matrix_shape( + feed_forward_dim = _matrix_shape(tensors, "action_model.model.transformer_blocks.0.ff.net.0.proj.weight" )[0] - mlp_hidden_dim, state_dim = matrix_shape("action_model.state_encoder.layer1.weight") - action_dim = matrix_shape("action_model.action_encoder.layer1.weight")[1] - future_token_count = matrix_shape("action_model.future_tokens.weight")[0] - max_sequence_length = matrix_shape("action_model.position_embedding.weight")[0] + mlp_hidden_dim, state_dim = _matrix_shape(tensors, "action_model.state_encoder.layer1.weight") + action_dim = _matrix_shape(tensors, "action_model.action_encoder.layer1.weight")[1] + future_token_count = _matrix_shape(tensors, "action_model.future_tokens.weight")[0] + max_sequence_length = _matrix_shape(tensors, "action_model.position_embedding.weight")[0] expected_shapes = { "action_model.model.timestep_encoder.timestep_embedder.linear_1.weight": [ @@ -676,32 +585,15 @@ def matrix_shape(name: str) -> list[int]: "action_model.position_embedding.weight": [max_sequence_length, dit_width], } for block in range(PI_BLOCK_COUNT): - prefix = f"action_model.model.transformer_blocks.{block}" expected_shapes.update( - { - f"{prefix}.norm1.linear.weight": [2 * dit_width, dit_width], - f"{prefix}.norm1.linear.bias": [2 * dit_width], - f"{prefix}.attn1.to_q.weight": [dit_width, dit_width], - f"{prefix}.attn1.to_q.bias": [dit_width], - f"{prefix}.attn1.to_k.weight": [dit_width, cross_attention_dim], - f"{prefix}.attn1.to_k.bias": [dit_width], - f"{prefix}.attn1.to_v.weight": [dit_width, cross_attention_dim], - f"{prefix}.attn1.to_v.bias": [dit_width], - f"{prefix}.attn1.to_out.0.weight": [dit_width, dit_width], - f"{prefix}.attn1.to_out.0.bias": [dit_width], - f"{prefix}.ff.net.0.proj.weight": [feed_forward_dim, dit_width], - f"{prefix}.ff.net.0.proj.bias": [feed_forward_dim], - f"{prefix}.ff.net.2.weight": [dit_width, feed_forward_dim], - f"{prefix}.ff.net.2.bias": [dit_width], - } + _dit_block_shapes( + f"action_model.model.transformer_blocks.{block}", + dit_width, + cross_attention_dim, + feed_forward_dim, + ) ) - mismatches = [ - f"{name}: expected {expected_shape}, got {shape(name)}" - for name, expected_shape in expected_shapes.items() - if shape(name) != expected_shape - ] - if mismatches: - raise StarVLAError("invalid legacy PI tensor shapes: " + "; ".join(mismatches)) + _validate_tensor_shapes(tensors, expected_shapes, label="legacy PI") return { "qwen_hidden_dim": cross_attention_dim, @@ -720,35 +612,26 @@ def matrix_shape(name: str) -> list[int]: def validate_pi_v3_tensors(tensors: dict[str, Any]) -> dict[str, int]: - """Validate every released PI_v3 policy tensor and infer its architecture.""" + """Validate PI_v3 policy tensors and infer their dimensions.""" actual = set(tensors) expected = set(PI_V3_TENSOR_MAP) missing = sorted(expected - actual) if missing: raise StarVLAError(f"PI-v3 policy is missing runtime tensors: {missing}") - def shape(name: str) -> list[int]: - return [int(dim) for dim in tensors[name].shape] - - def matrix_shape(name: str) -> list[int]: - value = shape(name) - if len(value) != 2: - raise StarVLAError(f"invalid PI_v3 matrix shape for {name}: {value}") - return value - - timestep_input = matrix_shape( + timestep_input = _matrix_shape(tensors, "action_model.model.timestep_encoder.timestep_embedder.linear_1.weight" ) dit_width, timestep_dim = timestep_input - feed_forward_dim = matrix_shape( + feed_forward_dim = _matrix_shape(tensors, "action_model.model.transformer_blocks.0.ff.net.0.proj.weight" )[0] - mlp_hidden_dim = matrix_shape("action_model.action_decoder.layer1.weight")[0] - action_dim = matrix_shape("action_model.action_encoder.layer1.weight")[1] - future_token_count = matrix_shape("action_model.future_tokens.weight")[0] - max_sequence_length = matrix_shape("action_model.position_embedding.weight")[0] - qwen_hidden_dim = shape("project_layers.0.0.weight")[0] - projector_output_dim = matrix_shape("project_layers.0.1.weight")[0] + mlp_hidden_dim = _matrix_shape(tensors, "action_model.action_decoder.layer1.weight")[0] + action_dim = _matrix_shape(tensors, "action_model.action_encoder.layer1.weight")[1] + future_token_count = _matrix_shape(tensors, "action_model.future_tokens.weight")[0] + max_sequence_length = _matrix_shape(tensors, "action_model.position_embedding.weight")[0] + qwen_hidden_dim = _tensor_shape(tensors, "project_layers.0.0.weight")[0] + projector_output_dim = _matrix_shape(tensors, "project_layers.0.1.weight")[0] if projector_output_dim != dit_width: raise StarVLAError( "invalid PI_v3 projector/DiT width contract: " @@ -780,24 +663,13 @@ def matrix_shape(name: str) -> list[int]: "action_model.position_embedding.weight": [max_sequence_length, dit_width], } for block in range(PI_V3_BLOCK_COUNT): - prefix = f"action_model.model.transformer_blocks.{block}" expected_shapes.update( - { - f"{prefix}.norm1.linear.weight": [2 * dit_width, dit_width], - f"{prefix}.norm1.linear.bias": [2 * dit_width], - f"{prefix}.attn1.to_q.weight": [dit_width, dit_width], - f"{prefix}.attn1.to_q.bias": [dit_width], - f"{prefix}.attn1.to_k.weight": [dit_width, dit_width], - f"{prefix}.attn1.to_k.bias": [dit_width], - f"{prefix}.attn1.to_v.weight": [dit_width, dit_width], - f"{prefix}.attn1.to_v.bias": [dit_width], - f"{prefix}.attn1.to_out.0.weight": [dit_width, dit_width], - f"{prefix}.attn1.to_out.0.bias": [dit_width], - f"{prefix}.ff.net.0.proj.weight": [feed_forward_dim, dit_width], - f"{prefix}.ff.net.0.proj.bias": [feed_forward_dim], - f"{prefix}.ff.net.2.weight": [dit_width, feed_forward_dim], - f"{prefix}.ff.net.2.bias": [dit_width], - } + _dit_block_shapes( + f"action_model.model.transformer_blocks.{block}", + dit_width, + dit_width, + feed_forward_dim, + ) ) for projector in range(PI_V3_PROJECTOR_COUNT): prefix = f"project_layers.{projector}" @@ -809,13 +681,7 @@ def matrix_shape(name: str) -> list[int]: f"{prefix}.1.bias": [dit_width], } ) - mismatches = [ - f"{name}: expected {expected_shape}, got {shape(name)}" - for name, expected_shape in expected_shapes.items() - if shape(name) != expected_shape - ] - if mismatches: - raise StarVLAError("invalid PI_v3 tensor shapes: " + "; ".join(mismatches)) + _validate_tensor_shapes(tensors, expected_shapes, label="PI_v3") return { "qwen_hidden_dim": qwen_hidden_dim, @@ -836,14 +702,18 @@ def load_variant_config( policy_dir: Path, surgery_manifest: dict[str, Any], variant_name: str, + backbone: str | None = None, ) -> dict[str, Any]: catalog_variant = str(surgery_manifest.get("variant", variant_name)) + effective_backbone = backbone or surgery_manifest.get("backbone") + if not isinstance(effective_backbone, str): + raise StarVLAError("surgery manifest does not identify the Qwen backbone") effective = resolve_effective_config( policy_dir, catalog_variant, { "framework": surgery_manifest.get("framework", variant_name), - "backbone": surgery_manifest.get("backbone", "qwen3_vl"), + "backbone": effective_backbone, }, ) effective_path = policy_dir / "effective_config.json" @@ -860,38 +730,34 @@ def load_variant_config( ) stored_effective = _load_json(effective_path) if stored_effective != effective: - # Qwen3 bundles produced before Qwen2.5 support predate these two - # explicit annotations. Qwen3 was the only supported backbone then, so - # this is an unambiguous legacy spelling of the same effective config. - legacy_effective = copy.deepcopy(effective) - legacy_metadata = legacy_effective.get("_robotcpp_effective_config") - if ( - surgery_manifest.get("backbone", "qwen3_vl") == "qwen3_vl" - and isinstance(legacy_metadata, dict) - ): - legacy_metadata.pop("backbone", None) - legacy_metadata.pop("framework", None) - if stored_effective != legacy_effective: - raise StarVLAError( - f"effective {variant_name.upper()} config does not match its canonical source/manifest" - ) + raise StarVLAError( + f"effective {variant_name.upper()} config does not match its canonical source/manifest" + ) return effective -def load_oft_config(policy_dir: Path, surgery_manifest: dict[str, Any]) -> dict[str, Any]: - return load_variant_config(policy_dir, surgery_manifest, "oft") +def load_oft_config( + policy_dir: Path, surgery_manifest: dict[str, Any], backbone: str +) -> dict[str, Any]: + return load_variant_config(policy_dir, surgery_manifest, "oft", backbone) -def load_groot_config(policy_dir: Path, surgery_manifest: dict[str, Any]) -> dict[str, Any]: - return load_variant_config(policy_dir, surgery_manifest, "groot") +def load_groot_config( + policy_dir: Path, surgery_manifest: dict[str, Any], backbone: str +) -> dict[str, Any]: + return load_variant_config(policy_dir, surgery_manifest, "groot", backbone) -def load_pi_config(policy_dir: Path, surgery_manifest: dict[str, Any]) -> dict[str, Any]: - return load_variant_config(policy_dir, surgery_manifest, "pi") +def load_pi_config( + policy_dir: Path, surgery_manifest: dict[str, Any], backbone: str +) -> dict[str, Any]: + return load_variant_config(policy_dir, surgery_manifest, "pi", backbone) -def load_pi_v3_config(policy_dir: Path, surgery_manifest: dict[str, Any]) -> dict[str, Any]: - return load_variant_config(policy_dir, surgery_manifest, "pi_v3") +def load_pi_v3_config( + policy_dir: Path, surgery_manifest: dict[str, Any], backbone: str +) -> dict[str, Any]: + return load_variant_config(policy_dir, surgery_manifest, "pi_v3", backbone) def resolve_action_token_id(hf_dir: Path) -> int: @@ -912,15 +778,22 @@ def resolve_action_token_id(hf_dir: Path) -> int: return token_ids[0] -def normalization_metadata(stats: dict[str, Any], action_dim: int) -> dict[str, Any]: +def normalization_metadata( + stats: dict[str, Any], action_dim: int, default_profile: str +) -> dict[str, Any]: + if default_profile not in stats: + raise StarVLAError( + f"default normalization profile {default_profile!r} is not present" + ) + profile_keys = [default_profile, *sorted(set(stats) - {default_profile})] metadata: dict[str, Any] = { "starvla.normalization.profile_count": len(stats), - "starvla.normalization.profile_keys": sorted(stats), + "starvla.normalization.profile_keys": profile_keys, "starvla.normalization.clip_actions": False, "starvla.normalization.binary_threshold": 0.5, "starvla.normalization.binary_comparison": "gt", } - for index, key in enumerate(sorted(stats)): + for index, key in enumerate(profile_keys): profile = stats[key] action = profile.get("action") if not isinstance(action, dict): @@ -989,17 +862,17 @@ def build_oft_metadata( text_filename: str, mmproj_filename: str, ) -> dict[str, Any]: - backbone = str(variant.get("backbone", "qwen3_vl")) - config = load_oft_config(policy_dir, surgery_manifest) + backbone = str(variant["backbone"]) + config = load_oft_config(policy_dir, surgery_manifest, backbone) framework = config.get("framework", {}) action_config = framework.get("action_model", {}) datasets = config.get("datasets", {}) vla_config = datasets.get("vla_data", {}) action_horizon = int(action_config.get("action_horizon", int(action_config.get("future_action_window_size", 15)) + 1)) if action_horizon != 16: - raise StarVLAError(f"unexpected official OFT action horizon: {action_horizon}") + raise StarVLAError(f"unsupported OFT action horizon: {action_horizon}") if dimensions["action_dim"] != 7: - raise StarVLAError(f"unexpected official OFT action dimension: {dimensions['action_dim']}") + raise StarVLAError(f"unsupported OFT action dimension: {dimensions['action_dim']}") expected_dimensions = { "qwen3_vl": (2560, 5120), "qwen2_5_vl": (2048, 4096), @@ -1010,14 +883,14 @@ def build_oft_metadata( dimensions["input_dim"], dimensions["hidden_dim"], ) != expected_dimensions: - raise StarVLAError(f"unexpected official OFT MLP dimensions: {dimensions}") + raise StarVLAError(f"unsupported OFT MLP dimensions: {dimensions}") action_tokens = OFT_ACTION_TOKEN * action_horizon action_suffix = f" Please predict the next {action_horizon} robot actions: {action_tokens}." image_size = vla_config.get("image_size", [224, 224]) image_names = vla_config.get("obs", ["image_0"]) - if image_size != [224, 224] or image_names != ["image_0"]: - raise StarVLAError(f"unexpected official OFT image contract: image_size={image_size}, obs={image_names}") + if image_size != [224, 224] or not isinstance(image_names, list): + raise StarVLAError(f"unsupported OFT image config: image_size={image_size}, obs={image_names}") qwen = ( _validate_pinned_qwen3vl_contract(hf_dir) @@ -1086,14 +959,11 @@ def build_oft_metadata( ) )) stats = _load_json(policy_dir / "dataset_statistics.json") - expected_profiles = ( - {"oxe_bridge", "oxe_rt1"} - if backbone == "qwen3_vl" - else {"bridge_dataset", "fractal20220817_data"} + metadata.update( + normalization_metadata( + stats, dimensions["action_dim"], variant["default_unnorm_key"] + ) ) - if set(stats) != expected_profiles: - raise StarVLAError(f"unexpected official OFT normalization profiles: {sorted(stats)}") - metadata.update(normalization_metadata(stats, dimensions["action_dim"])) return metadata @@ -1273,18 +1143,16 @@ def _validate_pinned_qwenvl_contract( raise StarVLAError(f"unsupported StarVLA Qwen backbone: {backbone!r}") -def _require_released_obs_pre_resize_disabled( +def _validate_image_config( vla_config: dict[str, Any], variant_label: str, config_label: str, ) -> None: if not isinstance(vla_config, dict): - raise StarVLAError(f"official {variant_label} {config_label} vla_data must be an object") + raise StarVLAError(f"{variant_label} {config_label} vla_data must be an object") if "obs_image_size" in vla_config: raise StarVLAError( - f"official {variant_label} {config_label} unexpectedly defines " - "datasets.vla_data.obs_image_size; released predict_action must leave its " - "optional pre-resize branch disabled" + f"{variant_label} does not support datasets.vla_data.obs_image_size" ) @@ -1296,10 +1164,10 @@ def build_qwen3vl_image_metadata( variant_label: str, config_label: str = "effective config", ) -> dict[str, Any]: - """Build the released dynamic Qwen3-VL image preprocessing contract.""" - _require_released_obs_pre_resize_disabled(vla_config, variant_label, config_label) - if image_names != ["image_0"]: - raise StarVLAError(f"unexpected official {variant_label} image names: {image_names!r}") + """Build the Qwen3-VL image preprocessing contract.""" + _validate_image_config(vla_config, variant_label, config_label) + if not image_names or any(not isinstance(name, str) or not name for name in image_names): + raise StarVLAError(f"{variant_label} image names must be non-empty strings") processor_size = qwen.get("processor_size") actual = { @@ -1339,10 +1207,13 @@ def build_qwen3vl_image_metadata( or QWEN3VL_PROCESSOR_MAX_PIXELS % token_area ): raise StarVLAError("internal Qwen3-VL smart-resize image-token bounds drift") - return { + metadata = { key: list(value) if isinstance(value, list) else value for key, value in QWEN3VL_DYNAMIC_IMAGE_METADATA.items() } + metadata["starvla.image.count"] = len(image_names) + metadata["starvla.image.names"] = list(image_names) + return metadata def build_qwen25vl_image_metadata( @@ -1354,13 +1225,11 @@ def build_qwen25vl_image_metadata( config_label: str = "effective config", ) -> dict[str, Any]: """Build the Transformers 4.57 fast Qwen2.5-VL image contract.""" - _require_released_obs_pre_resize_disabled( + _validate_image_config( vla_config, variant_label, config_label ) - if image_names != ["image_0"]: - raise StarVLAError( - f"unexpected official {variant_label} image names: {image_names!r}" - ) + if not image_names or any(not isinstance(name, str) or not name for name in image_names): + raise StarVLAError(f"{variant_label} image names must be non-empty strings") expected = { "processor_min_pixels": QWEN25VL_PROCESSOR_MIN_PIXELS, "processor_max_pixels": QWEN25VL_PROCESSOR_MAX_PIXELS, @@ -1394,7 +1263,7 @@ def build_qwen25vl_image_metadata( "internal Qwen2.5-VL smart-resize image-token bounds drift" ) return { - "starvla.image.count": 1, + "starvla.image.count": len(image_names), "starvla.image.names": list(image_names), "starvla.image.preprocessing_mode": "qwen2_5vl_smart_resize", "starvla.image.framework_inference_pre_resize": False, @@ -1446,20 +1315,20 @@ def build_groot_metadata( text_filename: str, mmproj_filename: str, ) -> dict[str, Any]: - """Build the executable contract for a released Qwen-VL GR00T head.""" - backbone = str(variant.get("backbone", "qwen3_vl")) - config = load_groot_config(policy_dir, surgery_manifest) + """Build the executable contract for a Qwen-VL GR00T head.""" + backbone = str(variant["backbone"]) + config = load_groot_config(policy_dir, surgery_manifest, backbone) framework = config.get("framework", {}) action_config = framework.get("action_model", {}) diffusion_config = action_config.get("diffusion_model_cfg", {}) vla_config = config.get("datasets", {}).get("vla_data", {}) qwen = _validate_pinned_qwenvl_contract(hf_dir, backbone) - expected_dimensions = GROOT_OFFICIAL_DIMENSIONS_BY_BACKBONE.get(backbone) + expected_dimensions = GROOT_SUPPORTED_DIMENSIONS_BY_BACKBONE.get(backbone) if expected_dimensions is None: raise StarVLAError(f"unsupported GR00T Qwen backbone: {backbone!r}") if dimensions != expected_dimensions: - raise StarVLAError(f"unexpected official GR00T tensor dimensions: {dimensions}") + raise StarVLAError(f"unsupported GR00T tensor dimensions: {dimensions}") if qwen["hidden_size"] != dimensions["qwen_hidden_dim"]: raise StarVLAError( "GR00T cross-attention dimension does not match the staged Qwen backbone" @@ -1471,102 +1340,26 @@ def build_groot_metadata( int(action_config.get("future_action_window_size", 15)) + 1, ) ) - expected_action_config = { - "action_model_type": action_config.get("action_model_type"), - "hidden_size": action_config.get("hidden_size"), - "add_pos_embed": action_config.get("add_pos_embed"), - "max_seq_len": action_config.get("max_seq_len"), - "action_dim": action_config.get("action_dim"), - "state_dim": action_config.get("state_dim"), - "action_horizon": action_horizon, - "past_action_window_size": action_config.get("past_action_window_size"), - "repeated_diffusion_steps": action_config.get("repeated_diffusion_steps"), - "noise_beta_alpha": action_config.get("noise_beta_alpha"), - "noise_beta_beta": action_config.get("noise_beta_beta"), - "noise_s": action_config.get("noise_s"), - "num_timestep_buckets": action_config.get("num_timestep_buckets"), - "num_inference_timesteps": action_config.get("num_inference_timesteps"), - "num_target_vision_tokens": action_config.get("num_target_vision_tokens"), - } - required_action_config = { - "action_model_type": "DiT-B", - "hidden_size": 1024, - "add_pos_embed": True, - "max_seq_len": 1024, - "action_dim": 7, - "state_dim": 7, - "action_horizon": 16, - "past_action_window_size": 0, - "repeated_diffusion_steps": 8, - "noise_beta_alpha": 1.5, - "noise_beta_beta": 1.0, - "noise_s": 0.999, - "num_timestep_buckets": 1000, - "num_inference_timesteps": 4, - "num_target_vision_tokens": 32, - } - if expected_action_config != required_action_config: - raise StarVLAError(f"unexpected official GR00T action config: {expected_action_config}") - - actual_diffusion_config = { - "input_embedding_dim": diffusion_config.get("input_embedding_dim"), - "attention_head_dim": diffusion_config.get("attention_head_dim"), - "num_attention_heads": diffusion_config.get("num_attention_heads"), - "cross_attention_dim": diffusion_config.get("cross_attention_dim"), - "dropout": diffusion_config.get("dropout"), - "final_dropout": diffusion_config.get("final_dropout"), - "interleave_self_attention": diffusion_config.get("interleave_self_attention"), - "norm_type": diffusion_config.get("norm_type"), - "num_layers": diffusion_config.get("num_layers"), - "output_dim": diffusion_config.get("output_dim"), - "positional_embeddings": diffusion_config.get("positional_embeddings"), - } - required_diffusion_config = { - "input_embedding_dim": 768, - "attention_head_dim": 64, - "num_attention_heads": 12, - "cross_attention_dim": dimensions["qwen_hidden_dim"], - "dropout": 0.2, - "final_dropout": True, - "interleave_self_attention": True, - "norm_type": "ada_norm", - "num_layers": 16, - "output_dim": 1024, - "positional_embeddings": None, - } - if actual_diffusion_config != required_diffusion_config: - raise StarVLAError(f"unexpected official GR00T diffusion config: {actual_diffusion_config}") - - framework_identity = ( - framework.get("name") - if backbone == "qwen3_vl" - else framework.get("framework_py") - ) - expected_framework_identity = ( - "QwenGR00T" if backbone == "qwen3_vl" else "QwenFM" - ) - if framework_identity != expected_framework_identity: - raise StarVLAError( - f"unexpected official GR00T framework identity: {framework_identity!r}" - ) - if vla_config.get("image_size") != [224, 224] or vla_config.get("obs") != ["image_0"]: - raise StarVLAError( - "unexpected official GR00T image contract: " - f"image_size={vla_config.get('image_size')}, obs={vla_config.get('obs')}" - ) + if action_horizon != 16: + raise StarVLAError(f"unsupported GR00T action horizon: {action_horizon}") + head_count = int(diffusion_config.get("num_attention_heads", 0)) + head_dim = int(diffusion_config.get("attention_head_dim", 0)) + if head_count * head_dim != dimensions["dit_width"]: + raise StarVLAError("GR00T attention config does not match checkpoint shapes") + image_names = vla_config.get("obs", ["image_0"]) + if vla_config.get("image_size", [224, 224]) != [224, 224] or not isinstance(image_names, list): + raise StarVLAError("GR00T requires a 224x224 image configuration") if vla_config.get("include_state", False) not in (False, "False"): - raise StarVLAError("released GR00T checkpoint unexpectedly enables training state input") + raise StarVLAError("GR00T state input is not supported") cot_template = str(vla_config.get("CoT_prompt", "")) - required_cot = ( - "Your task is {instruction}. To identify the key objects for your task. " - "Locate their bounding boxes in [x1,y1,x2,y2] format." - ) - if cot_template != required_cot: - raise StarVLAError(f"unexpected official GR00T CoT prompt: {cot_template!r}") + if not cot_template: + raise StarVLAError("GR00T CoT_prompt must not be empty") num_steps = int(action_config["num_inference_timesteps"]) timestep_buckets = int(action_config["num_timestep_buckets"]) + if num_steps <= 0 or timestep_buckets <= 0: + raise StarVLAError("GR00T timestep counts must be positive") timestep_ids = [step * timestep_buckets // num_steps for step in range(num_steps)] metadata: dict[str, Any] = { "general.architecture": "starvla-policy", @@ -1596,8 +1389,8 @@ def build_groot_metadata( "starvla.action.binary_dimensions": [6], "starvla.groot.dit_width": dimensions["dit_width"], "starvla.groot.block_count": dimensions["block_count"], - "starvla.groot.attention_head_count": 12, - "starvla.groot.attention_head_dim": 64, + "starvla.groot.attention_head_count": head_count, + "starvla.groot.attention_head_dim": head_dim, "starvla.groot.cross_attention_dim": dimensions["qwen_hidden_dim"], "starvla.groot.feed_forward_dim": dimensions["feed_forward_dim"], "starvla.groot.ada_norm_epsilon": GROOT_DIT_NORM_EPS, @@ -1616,7 +1409,7 @@ def build_groot_metadata( build_qwen3vl_image_metadata( vla_config, qwen, - ["image_0"], + image_names, variant_label="GR00T", ) )) @@ -1625,22 +1418,16 @@ def build_groot_metadata( build_qwen25vl_image_metadata( vla_config, qwen, - ["image_0"], + image_names, variant_label="GR00T", ) )) stats = _load_json(policy_dir / "dataset_statistics.json") - if set(stats) != {"oxe_bridge", "oxe_rt1"}: - raise StarVLAError(f"unexpected official GR00T normalization profiles: {sorted(stats)}") - state_dimensions = sorted( - { - len(profile.get("state", {}).get("q01", [])) - for profile in stats.values() - } + metadata.update( + normalization_metadata( + stats, dimensions["action_dim"], variant["default_unnorm_key"] + ) ) - if state_dimensions != [8]: - raise StarVLAError(f"unexpected official GR00T state statistics dimensions: {state_dimensions}") - metadata.update(normalization_metadata(stats, dimensions["action_dim"])) return metadata @@ -1653,128 +1440,52 @@ def build_pi_metadata( text_filename: str, mmproj_filename: str, ) -> dict[str, Any]: - """Build the released Qwen2.5-VL legacy PI executable contract.""" + """Build the Qwen2.5-VL legacy PI executable contract.""" if variant.get("framework") != "pi" or variant.get("backbone") != "qwen2_5_vl": raise StarVLAError("legacy PI metadata requires the qwen25_pi catalog variant") - config = load_pi_config(policy_dir, surgery_manifest) + config = load_pi_config(policy_dir, surgery_manifest, str(variant["backbone"])) framework = config.get("framework", {}) action_config = framework.get("action_model", {}) diffusion_config = action_config.get("diffusion_model_cfg", {}) vla_config = config.get("datasets", {}).get("vla_data", {}) qwen = _validate_pinned_qwen25vl_contract(hf_dir) - if dimensions != PI_OFFICIAL_DIMENSIONS: - raise StarVLAError(f"unexpected official legacy PI tensor dimensions: {dimensions}") + if dimensions != PI_SUPPORTED_DIMENSIONS: + raise StarVLAError(f"unsupported legacy PI tensor dimensions: {dimensions}") if qwen["hidden_size"] != dimensions["qwen_hidden_dim"]: raise StarVLAError( "legacy PI cross-attention dimension does not match the staged Qwen backbone" ) - actual_action_config = { - "action_model_type": action_config.get("action_model_type"), - "hidden_size": action_config.get("hidden_size"), - "action_hidden_dim": action_config.get("action_hidden_dim"), - "add_pos_embed": action_config.get("add_pos_embed"), - "max_seq_len": action_config.get("max_seq_len"), - "action_dim": action_config.get("action_dim"), - "state_dim": action_config.get("state_dim"), - "future_action_window_size": action_config.get("future_action_window_size"), - "action_horizon": action_config.get("action_horizon"), - "past_action_window_size": action_config.get("past_action_window_size"), - "repeated_diffusion_steps": action_config.get("repeated_diffusion_steps"), - "noise_beta_alpha": action_config.get("noise_beta_alpha"), - "noise_beta_beta": action_config.get("noise_beta_beta"), - "noise_s": action_config.get("noise_s"), - "num_timestep_buckets": action_config.get("num_timestep_buckets"), - "num_inference_timesteps": action_config.get("num_inference_timesteps"), - "num_target_vision_tokens": action_config.get("num_target_vision_tokens"), - } - required_action_config = { - "action_model_type": "DiT-Qwen", - "hidden_size": 2048, - "action_hidden_dim": 2048, - "add_pos_embed": True, - "max_seq_len": 1024, - "action_dim": 7, - "state_dim": 7, - "future_action_window_size": 15, - "action_horizon": 16, - "past_action_window_size": 0, - "repeated_diffusion_steps": 8, - "noise_beta_alpha": 1.5, - "noise_beta_beta": 1.0, - "noise_s": 0.999, - "num_timestep_buckets": 1000, - "num_inference_timesteps": 4, - "num_target_vision_tokens": 32, - } - if actual_action_config != required_action_config: - raise StarVLAError( - f"unexpected official legacy PI action config: {actual_action_config}" - ) - - actual_diffusion_config = { - "input_embedding_dim": diffusion_config.get("input_embedding_dim"), - "attention_head_dim": diffusion_config.get("attention_head_dim"), - "num_attention_heads": diffusion_config.get("num_attention_heads"), - "cross_attention_dim": diffusion_config.get("cross_attention_dim"), - "dropout": diffusion_config.get("dropout"), - "final_dropout": diffusion_config.get("final_dropout"), - "interleave_self_attention": diffusion_config.get("interleave_self_attention"), - "use_canonical_forward": diffusion_config.get("use_canonical_forward"), - "norm_type": diffusion_config.get("norm_type"), - "num_layers": diffusion_config.get("num_layers"), - "output_dim": diffusion_config.get("output_dim"), - "positional_embeddings": diffusion_config.get("positional_embeddings"), - } - required_diffusion_config = { - "input_embedding_dim": 2048, - "attention_head_dim": 64, - "num_attention_heads": 32, - "cross_attention_dim": 2048, - "dropout": 0.2, - "final_dropout": True, - "interleave_self_attention": True, - "use_canonical_forward": False, - "norm_type": "ada_norm", - "num_layers": 16, - "output_dim": 1024, - "positional_embeddings": None, - } - if actual_diffusion_config != required_diffusion_config: - raise StarVLAError( - f"unexpected official legacy PI diffusion config: {actual_diffusion_config}" - ) - if framework.get("name") != "QwenPI": - raise StarVLAError( - f"unexpected official legacy PI framework name: {framework.get('name')!r}" - ) - qwen_config = framework.get("qwenvl", {}) - if ( - qwen_config.get("vl_hidden_dim") != dimensions["qwen_hidden_dim"] - or qwen_config.get("attn_implementation") != "flash_attention_2" - ): - raise StarVLAError( - f"unexpected official legacy PI Qwen contract: {qwen_config}" + action_horizon = int( + action_config.get( + "action_horizon", + int(action_config.get("future_action_window_size", 15)) + 1, ) - - required_cot = ( - "Your task is {instruction}. To identify the key objects for your task. " - "Locate their bounding boxes in [x1,y1,x2,y2] format." ) + if action_horizon != 16: + raise StarVLAError(f"unsupported PI action horizon: {action_horizon}") + head_count = int(diffusion_config.get("num_attention_heads", 0)) + head_dim = int(diffusion_config.get("attention_head_dim", 0)) + if head_count * head_dim != dimensions["dit_width"]: + raise StarVLAError("PI attention config does not match checkpoint shapes") cot_template = str(vla_config.get("CoT_prompt", "")) + image_names = vla_config.get("obs", ["image_0"]) + image_size = vla_config.get("image_size", [224, 224]) + if not cot_template: + raise StarVLAError("PI CoT_prompt must not be empty") if ( - cot_template != required_cot - or vla_config.get("obs") != ["image_0"] - or vla_config.get("image_size") != [224, 224] - or vla_config.get("default_image_resolution") != [3, 224, 224] - or vla_config.get("data_mix") != "bridge_rt_1" - or vla_config.get("action_type") != "delta_ee" + not isinstance(image_names, list) + or not isinstance(image_size, list) + or len(image_size) != 2 + or any(type(value) is not int or value <= 0 for value in image_size) ): - raise StarVLAError(f"unexpected official legacy PI VLA config: {vla_config}") + raise StarVLAError("PI image configuration is invalid") num_steps = int(action_config["num_inference_timesteps"]) timestep_buckets = int(action_config["num_timestep_buckets"]) + if num_steps <= 0 or timestep_buckets <= 0: + raise StarVLAError("PI timestep counts must be positive") continuous_times = [step / float(num_steps) for step in range(num_steps)] timestep_ids = [int(value * timestep_buckets) for value in continuous_times] hidden_tuple_indices = list( @@ -1798,14 +1509,14 @@ def build_pi_metadata( "starvla.prompt.cot_template": cot_template, "starvla.conditioning.hidden_tuple_indices": hidden_tuple_indices, "starvla.action.dimension": dimensions["action_dim"], - "starvla.action.horizon": 16, + "starvla.action.horizon": action_horizon, "starvla.action.continuous_dimensions": [0, 1, 2, 3, 4, 5], "starvla.action.binary_dimensions": [6], "starvla.state.dimension": dimensions["state_dim"], "starvla.pi.dit_width": dimensions["dit_width"], "starvla.pi.block_count": dimensions["block_count"], - "starvla.pi.attention_head_count": 32, - "starvla.pi.attention_head_dim": 64, + "starvla.pi.attention_head_count": head_count, + "starvla.pi.attention_head_dim": head_dim, "starvla.pi.cross_attention_dim": dimensions["qwen_hidden_dim"], "starvla.pi.feed_forward_dim": dimensions["feed_forward_dim"], "starvla.pi.mlp_hidden_dimension": dimensions["mlp_hidden_dim"], @@ -1821,7 +1532,7 @@ def build_pi_metadata( image_metadata = build_qwen25vl_image_metadata( vla_config, qwen, - ["image_0"], + image_names, variant_label="legacy PI", ) image_metadata.update( @@ -1829,8 +1540,8 @@ def build_pi_metadata( "starvla.image.framework_inference_pre_resize": True, "starvla.image.framework_inference_pre_resize_config_key": "datasets.vla_data.image_size", - "starvla.image.framework_inference_pre_resize_width": 224, - "starvla.image.framework_inference_pre_resize_height": 224, + "starvla.image.framework_inference_pre_resize_width": image_size[1], + "starvla.image.framework_inference_pre_resize_height": image_size[0], } ) for key in ( @@ -1848,21 +1559,11 @@ def build_pi_metadata( metadata[key] = image_metadata[key] stats = _load_json(policy_dir / "dataset_statistics.json") - if set(stats) != {"oxe_bridge", "oxe_rt1"}: - raise StarVLAError( - f"unexpected official legacy PI normalization profiles: {sorted(stats)}" + metadata.update( + normalization_metadata( + stats, dimensions["action_dim"], variant["default_unnorm_key"] ) - state_dimensions = sorted( - { - len(profile.get("state", {}).get("q01", [])) - for profile in stats.values() - } ) - if state_dimensions != [8]: - raise StarVLAError( - f"unexpected official legacy PI state statistics dimensions: {state_dimensions}" - ) - metadata.update(normalization_metadata(stats, dimensions["action_dim"])) metadata["starvla.normalization.clip_actions"] = True metadata["starvla.normalization.binary_comparison"] = "ge" return metadata @@ -1877,13 +1578,14 @@ def build_pi_v3_metadata( text_filename: str, mmproj_filename: str, ) -> dict[str, Any]: - config = load_pi_v3_config(policy_dir, surgery_manifest) - full_config = _load_yaml(policy_dir / "config.full.yaml") + config = load_pi_v3_config( + policy_dir, surgery_manifest, str(variant["backbone"]) + ) framework = config.get("framework", {}) action = framework.get("action_model", {}) diffusion = action.get("diffusion_model_cfg", {}) vla = config.get("datasets", {}).get("vla_data", {}) - image_names = full_config.get("datasets", {}).get("vla_data", {}).get("obs") + image_names = vla.get("obs", ["image_0"]) qwen = _validate_pinned_qwen3vl_contract(hf_dir) expected_dimensions = { @@ -1892,16 +1594,22 @@ def build_pi_v3_metadata( "action_dim": action.get("action_dim"), "block_count": diffusion.get("num_layers"), } - if framework.get("name") != "QwenPI_v3" or any( - dimensions[key] != value for key, value in expected_dimensions.items() - ): + if any(dimensions[key] != value for key, value in expected_dimensions.items()): raise StarVLAError("PI-v3 config does not match the checkpoint tensor shapes") - if not isinstance(image_names, list) or not image_names: + if not isinstance(image_names, list) or not image_names or any( + not isinstance(name, str) or not name for name in image_names + ): raise StarVLAError("PI-v3 config does not define observation image names") horizon = int(action["action_horizon"]) num_steps = int(action["num_inference_timesteps"]) timestep_buckets = int(action["num_timestep_buckets"]) + if horizon != 16: + raise StarVLAError(f"unsupported PI-v3 action horizon: {horizon}") + if num_steps <= 0 or timestep_buckets <= 0: + raise StarVLAError("PI-v3 timestep counts must be positive") + if not str(vla.get("CoT_prompt", "")): + raise StarVLAError("PI-v3 CoT_prompt must not be empty") processor_size = qwen["processor_size"] metadata: dict[str, Any] = { "general.architecture": "starvla-policy", @@ -1952,12 +1660,13 @@ def build_pi_v3_metadata( normalization_metadata( _load_json(policy_dir / "dataset_statistics.json"), dimensions["action_dim"], + variant["default_unnorm_key"], ) ) return metadata -def convert_oft_policy( +def convert_policy( policy_dir: Path, hf_dir: Path, surgery_manifest_path: Path, @@ -1970,213 +1679,83 @@ def convert_oft_policy( catalog = load_catalog(catalog_path) surgery_manifest = _load_json(surgery_manifest_path) variant = get_variant(catalog, str(surgery_manifest.get("variant", ""))) - if variant.get("framework") != "oft": + framework = str(variant["framework"]) + validators = { + "oft": validate_oft_tensors, + "groot": validate_groot_tensors, + "pi": validate_pi_tensors, + "pi_v3": validate_pi_v3_tensors, + } + tensor_maps = { + "oft": OFT_TENSOR_MAP, + "groot": GROOT_TENSOR_MAP, + "pi": PI_TENSOR_MAP, + "pi_v3": PI_V3_TENSOR_MAP, + } + if framework not in validators: raise StarVLAError( - f"surgery variant {surgery_manifest.get('variant')!r} is not an OFT policy" + f"surgery variant {surgery_manifest.get('variant')!r} has no policy converter" ) - validate_official_surgery_manifest(surgery_manifest, variant, catalog) - verify_staged_assets(hf_dir, surgery_manifest.get("qwen_assets", {}), component="Qwen") - verify_staged_assets(policy_dir, surgery_manifest.get("policy_assets", {}), component="policy") - verify_staged_tensors_against_checkpoint( - policy_dir, - surgery_manifest.get("policy_output", {}), - surgery_manifest, - variant, - component="policy", - ) - tensors = load_policy_tensors(policy_dir) - dimensions = validate_oft_tensors(tensors) - action_token_id = resolve_action_token_id(hf_dir) - metadata = build_oft_metadata( - policy_dir, - hf_dir, - variant, - surgery_manifest, - dimensions, - action_token_id, - text_filename, - mmproj_filename, + validate_surgery_manifest(surgery_manifest, variant, catalog) + verify_staged_assets( + hf_dir, surgery_manifest.get("qwen_assets", {}), component="Qwen" ) - - pi0_writer_dir = Path(__file__).resolve().parents[1] / "pi0" - sys.path.insert(0, str(pi0_writer_dir)) - try: - from gguf_writer import write_gguf_arrays - except ImportError as exc: - raise StarVLAError(f"failed to import repository GGUF writer adapter: {exc}") from exc - - def arrays(): - for source_name, destination_name in OFT_TENSOR_MAP.items(): - tensor = tensors[source_name] - array = tensor.detach().float().cpu().numpy() - yield destination_name, [int(dim) for dim in tensor.shape], np.asarray(array), dtype - - _write_gguf_arrays_no_overwrite(output, metadata, arrays(), write_gguf_arrays) - - -def convert_groot_policy( - policy_dir: Path, - hf_dir: Path, - surgery_manifest_path: Path, - output: Path, - catalog_path: Path, - dtype: str, - text_filename: str, - mmproj_filename: str, -) -> None: - catalog = load_catalog(catalog_path) - surgery_manifest = _load_json(surgery_manifest_path) - variant = get_variant(catalog, str(surgery_manifest.get("variant", ""))) - if variant.get("framework") != "groot": - raise StarVLAError( - f"surgery variant {surgery_manifest.get('variant')!r} is not a GR00T policy" - ) - validate_official_surgery_manifest(surgery_manifest, variant, catalog) - verify_staged_assets(hf_dir, surgery_manifest.get("qwen_assets", {}), component="Qwen") - verify_staged_assets(policy_dir, surgery_manifest.get("policy_assets", {}), component="policy") - verify_staged_tensors_against_checkpoint( - policy_dir, - surgery_manifest.get("policy_output", {}), - surgery_manifest, - variant, - component="policy", + verify_staged_assets( + policy_dir, surgery_manifest.get("policy_assets", {}), component="policy" ) - tensors = load_policy_tensors(policy_dir) - dimensions = validate_groot_tensors(tensors) - metadata = build_groot_metadata( + dimensions = validators[framework](tensors) + metadata_args = ( policy_dir, hf_dir, variant, surgery_manifest, dimensions, - text_filename, - mmproj_filename, ) - - pi0_writer_dir = Path(__file__).resolve().parents[1] / "pi0" - sys.path.insert(0, str(pi0_writer_dir)) - try: - from gguf_writer import write_gguf_arrays - except ImportError as exc: - raise StarVLAError(f"failed to import repository GGUF writer adapter: {exc}") from exc - - def arrays(): - for source_name, destination_name in GROOT_TENSOR_MAP.items(): - tensor = tensors[source_name] - array = tensor.detach().float().cpu().numpy() - yield destination_name, [int(dim) for dim in tensor.shape], np.asarray(array), dtype - - _write_gguf_arrays_no_overwrite(output, metadata, arrays(), write_gguf_arrays) - - -def convert_pi_policy( - policy_dir: Path, - hf_dir: Path, - surgery_manifest_path: Path, - output: Path, - catalog_path: Path, - dtype: str, - text_filename: str, - mmproj_filename: str, -) -> None: - catalog = load_catalog(catalog_path) - surgery_manifest = _load_json(surgery_manifest_path) - variant = get_variant(catalog, str(surgery_manifest.get("variant", ""))) - if variant.get("framework") != "pi" or variant.get("backbone") != "qwen2_5_vl": - raise StarVLAError( - f"surgery variant {surgery_manifest.get('variant')!r} is not a Qwen2.5 legacy PI policy" + if framework == "oft": + metadata = build_oft_metadata( + *metadata_args, + resolve_action_token_id(hf_dir), + text_filename, + mmproj_filename, + ) + else: + metadata_builder = { + "groot": build_groot_metadata, + "pi": build_pi_metadata, + "pi_v3": build_pi_v3_metadata, + }[framework] + metadata = metadata_builder( + *metadata_args, + text_filename, + mmproj_filename, ) - validate_official_surgery_manifest(surgery_manifest, variant, catalog) - verify_staged_assets(hf_dir, surgery_manifest.get("qwen_assets", {}), component="Qwen") - verify_staged_assets(policy_dir, surgery_manifest.get("policy_assets", {}), component="policy") - verify_staged_tensors_against_checkpoint( - policy_dir, - surgery_manifest.get("policy_output", {}), - surgery_manifest, - variant, - component="policy", - ) - - tensors = load_policy_tensors(policy_dir) - dimensions = validate_pi_tensors(tensors) - metadata = build_pi_metadata( - policy_dir, - hf_dir, - variant, - surgery_manifest, - dimensions, - text_filename, - mmproj_filename, - ) - pi0_writer_dir = Path(__file__).resolve().parents[1] / "pi0" - sys.path.insert(0, str(pi0_writer_dir)) + writer_dir = Path(__file__).resolve().parents[1] / "pi0" + sys.path.insert(0, str(writer_dir)) try: from gguf_writer import write_gguf_arrays except ImportError as exc: - raise StarVLAError(f"failed to import repository GGUF writer adapter: {exc}") from exc + raise StarVLAError( + f"failed to import repository GGUF writer adapter: {exc}" + ) from exc def arrays(): - for source_name, destination_name in PI_TENSOR_MAP.items(): + for source_name, destination_name in tensor_maps[framework].items(): tensor = tensors[source_name] array = tensor.detach().float().cpu().numpy() - yield destination_name, [int(dim) for dim in tensor.shape], np.asarray(array), dtype - - _write_gguf_arrays_no_overwrite(output, metadata, arrays(), write_gguf_arrays) - - -def convert_pi_v3_policy( - policy_dir: Path, - hf_dir: Path, - surgery_manifest_path: Path, - output: Path, - catalog_path: Path, - dtype: str, - text_filename: str, - mmproj_filename: str, -) -> None: - catalog = load_catalog(catalog_path) - variant = get_variant(catalog, "pi_v3") - surgery_manifest = _load_json(surgery_manifest_path) - validate_official_surgery_manifest(surgery_manifest, variant, catalog) - verify_staged_assets(hf_dir, surgery_manifest.get("qwen_assets", {}), component="Qwen") - verify_staged_assets(policy_dir, surgery_manifest.get("policy_assets", {}), component="policy") - verify_staged_tensors_against_checkpoint( - policy_dir, - surgery_manifest.get("policy_output", {}), - surgery_manifest, - variant, - component="policy", - ) + yield ( + destination_name, + [int(dimension) for dimension in tensor.shape], + np.asarray(array), + dtype, + ) - tensors = load_policy_tensors(policy_dir) - dimensions = validate_pi_v3_tensors(tensors) - metadata = build_pi_v3_metadata( - policy_dir, - hf_dir, - variant, - surgery_manifest, - dimensions, - text_filename, - mmproj_filename, + _write_gguf_arrays_no_overwrite( + output, metadata, arrays(), write_gguf_arrays ) - pi0_writer_dir = Path(__file__).resolve().parents[1] / "pi0" - sys.path.insert(0, str(pi0_writer_dir)) - try: - from gguf_writer import write_gguf_arrays - except ImportError as exc: - raise StarVLAError(f"failed to import repository GGUF writer adapter: {exc}") from exc - - def arrays(): - for source_name, destination_name in PI_V3_TENSOR_MAP.items(): - tensor = tensors[source_name] - array = tensor.detach().float().cpu().numpy() - yield destination_name, [int(dim) for dim in tensor.shape], np.asarray(array), dtype - - _write_gguf_arrays_no_overwrite(output, metadata, arrays(), write_gguf_arrays) def parse_args() -> argparse.Namespace: @@ -2219,16 +1798,7 @@ def main() -> int: mmproj_filename = args.mmproj_filename or default_mmproj_filename( args.variant, DEFAULT_MMPROJ_DTYPE ) - converters = { - "oft": convert_oft_policy, - "groot": convert_groot_policy, - "pi_v3": convert_pi_v3_policy, - "qwen25_oft": convert_oft_policy, - "qwen25_groot": convert_groot_policy, - "qwen25_pi": convert_pi_policy, - } - converter = converters[args.variant] - converter( + convert_policy( policy_dir=args.policy_dir, hf_dir=args.hf_dir, surgery_manifest_path=args.surgery_manifest, diff --git a/tools/hf2gguf/starvla/convert_starvla_qwen25_fast.py b/tools/hf2gguf/starvla/convert_starvla_qwen25_fast.py index 8c17940..f9e4ef2 100644 --- a/tools/hf2gguf/starvla/convert_starvla_qwen25_fast.py +++ b/tools/hf2gguf/starvla/convert_starvla_qwen25_fast.py @@ -1,5 +1,5 @@ #!/usr/bin/env python3 -"""Stage and convert the official Qwen2.5-VL StarVLA FAST checkpoint.""" +"""Stage and convert a Qwen2.5-VL StarVLA FAST checkpoint.""" from __future__ import annotations @@ -15,6 +15,7 @@ import numpy as np +from convert_starvla_policy_to_gguf import normalization_metadata from convert_starvla_qwen_to_gguf import build_commands, verify_llama_checkout from starvla_checkpoint import ( DEFAULT_CATALOG, @@ -26,7 +27,8 @@ inventory_summary, load_catalog, load_checkpoint_state, - official_bundle_uuid, + bundle_uuid, + portable_source_record, sha256_file, staged_qwen_asset_hashes, validate_qwen_vlm_destination_names, @@ -105,17 +107,9 @@ MMPROJ_FILENAME = "mmproj-qwen25-fast-bf16.gguf" POLICY_FILENAME = "policy-qwen25-fast.gguf" STAGING_MANIFEST_FILENAME = "qwen25-fast-staging-manifest.json" -BUNDLE_MANIFEST_FILENAME = "qwen25-fast-bundle-manifest.json" +BUNDLE_MANIFEST_FILENAME = "conversion_manifest.json" -COT_PROMPT = ( - "Your task is {instruction}. To identify the key objects for your task. " - "Locate their bounding boxes in [x1,y1,x2,y2] format." -) ACTION_NAMES = ["x", "y", "z", "roll", "pitch", "yaw", "gripper"] -EXPECTED_NORMALIZATION_PROFILES = { - "bridge_dataset", - "fractal20220817_data", -} ACTION_TOKEN_MAP_TENSOR = "starvla.policy.fast.action_token_map" CODEC_TOKEN_OFFSETS_TENSOR = "starvla.policy.fast.codec.token_offsets" @@ -153,7 +147,6 @@ def validate_catalog_contract( "framework": FRAMEWORK, "backbone": BACKBONE, "model_type": MODEL_TYPE, - "status": "official_policy", "qwen_asset": QWEN_ASSET_KEY, "policy_prefixes": [], } @@ -162,18 +155,9 @@ def validate_catalog_contract( for key, value in expected.items() if entry.get(key) != value ] - checkpoint = entry.get("checkpoint") - if ( - not isinstance(checkpoint, Mapping) - or checkpoint.get("path") != "checkpoints/steps_10000_pytorch_model.pt" - or checkpoint.get("size") != 8_146_439_050 - or checkpoint.get("sha256") - != "f30e89a6b2a166fa3f48af42d5cffde07be44074b861abc7b57e1ccdb734e81e" - ): - mismatches.append("checkpoint: not the reviewed steps_10000 source lock") if entry.get("policy_tensors") not in (None, []): mismatches.append("policy_tensors: FAST must not split a separate policy head") - official_bundle_uuid(entry, catalog) + bundle_uuid(entry, catalog) qwen_name, qwen_entry = get_qwen_asset(catalog, entry) if qwen_name != QWEN_ASSET_KEY: @@ -557,88 +541,6 @@ def compile_fast_runtime_tensors( } -def normalization_metadata(stats: dict[str, Any], action_dim: int) -> dict[str, Any]: - if set(stats) != EXPECTED_NORMALIZATION_PROFILES: - raise StarVLAError( - "unexpected official FAST normalization profiles: " - f"{sorted(stats)}" - ) - metadata: dict[str, Any] = { - "starvla.normalization.profile_count": len(stats), - "starvla.normalization.profile_keys": sorted(stats), - "starvla.normalization.clip_actions": False, - "starvla.normalization.binary_threshold": 0.5, - "starvla.normalization.binary_comparison": "gt", - } - expected_mask = [True] * (action_dim - 1) + [False] - for index, key in enumerate(sorted(stats)): - profile = stats[key] - if not isinstance(profile, dict): - raise StarVLAError(f"normalization profile {key!r} must be an object") - action = profile.get("action") - if not isinstance(action, dict): - raise StarVLAError(f"normalization profile {key!r} has no action object") - for field in ("q01", "q99", "mask"): - values = action.get(field) - if not isinstance(values, list) or len(values) != action_dim: - raise StarVLAError( - f"normalization profile {key!r} action.{field} must " - f"have {action_dim} values" - ) - metadata[f"starvla.normalization.profile.{index}.action_{field}"] = values - q01 = action["q01"] - q99 = action["q99"] - mask = action["mask"] - if any(type(value) is not bool for value in mask) or mask != expected_mask: - raise StarVLAError( - f"normalization profile {key!r} action.mask must be {expected_mask}" - ) - if any( - isinstance(value, bool) - or not isinstance(value, (int, float)) - or not math.isfinite(value) - for value in [*q01, *q99] - ): - raise StarVLAError( - f"normalization profile {key!r} action quantiles must be finite" - ) - if any(q99[axis] < q01[axis] for axis in range(action_dim - 1)): - raise StarVLAError( - f"normalization profile {key!r} has q99 below q01" - ) - metadata[f"starvla.normalization.profile.{index}.key"] = key - - state = profile.get("state") - if not isinstance(state, dict): - raise StarVLAError( - f"normalization profile {key!r} has no state statistics" - ) - state_q01 = state.get("q01") - state_q99 = state.get("q99") - if ( - not isinstance(state_q01, list) - or not isinstance(state_q99, list) - or not state_q01 - or len(state_q01) != len(state_q99) - or any( - isinstance(value, bool) - or not isinstance(value, (int, float)) - or not math.isfinite(value) - for value in [*state_q01, *state_q99] - ) - or any(upper < lower for lower, upper in zip(state_q01, state_q99)) - ): - raise StarVLAError( - f"normalization profile {key!r} has invalid state q01/q99" - ) - metadata[f"starvla.normalization.profile.{index}.state_dimension"] = len( - state_q01 - ) - metadata[f"starvla.normalization.profile.{index}.state_q01"] = state_q01 - metadata[f"starvla.normalization.profile.{index}.state_q99"] = state_q99 - return metadata - - def _normalize_gguf_metadata_value(value: Any) -> Any: if isinstance(value, bool) or isinstance(value, str) or value is None: return value @@ -671,7 +573,6 @@ def build_fast_runtime_policy( qwen_dir: Path, codec_dir: Path, ) -> tuple[dict[str, Any], dict[str, np.ndarray]]: - del entry source = manifest.get("source") bundle_uuid = manifest.get("bundle_uuid") if not isinstance(source, Mapping) or not isinstance(bundle_uuid, str): @@ -687,8 +588,7 @@ def build_fast_runtime_policy( offsets = arrays[CODEC_TOKEN_OFFSETS_TENSOR] token_bytes = arrays[CODEC_TOKEN_BYTES_TENSOR] if ( - effective.get("cot_prompt") != COT_PROMPT - or effective.get("image_count") != 1 + effective.get("image_count") != 1 or effective.get("action_dim") != ACTION_DIM or effective.get("action_horizon") != ACTION_HORIZON ): @@ -715,7 +615,7 @@ def build_fast_runtime_policy( "starvla.action.continuous_dimensions": list(range(ACTION_DIM - 1)), "starvla.action.binary_dimensions": [ACTION_DIM - 1], "starvla.image.count": effective["image_count"], - "starvla.image.names": ["image_0"], + "starvla.image.names": effective["image_names"], "starvla.image.processor_min_pixels": processor["min_pixels"], "starvla.image.processor_max_pixels": processor["max_pixels"], "starvla.image.patch_size": processor["patch_size"], @@ -735,7 +635,9 @@ def build_fast_runtime_policy( "starvla.fast.codec.token_offsets_count": int(offsets.size), "starvla.fast.codec.token_bytes_count": int(token_bytes.size), } - metadata.update(normalization_metadata(stats, ACTION_DIM)) + metadata.update( + normalization_metadata(stats, ACTION_DIM, str(entry["default_unnorm_key"])) + ) return { key: _normalize_gguf_metadata_value(value) for key, value in metadata.items() @@ -935,6 +837,7 @@ def validate_fast_runtime_policy_gguf( def build_bundle_manifest( *, manifest: Mapping[str, Any], + entry: Mapping[str, Any], codec: Mapping[str, Any], text_component: Mapping[str, Any], mmproj_component: Mapping[str, Any], @@ -944,13 +847,13 @@ def build_bundle_manifest( raise StarVLAError("FAST policy component has an unexpected filename") return { "schema_version": 1, - "kind": "starvla_qwen25_fast_official_gguf_bundle", + "kind": "starvla_qwen25_fast_gguf_bundle", "variant": VARIANT_KEY, "framework": FRAMEWORK, "backbone": BACKBONE, "model_type": MODEL_TYPE, "bundle_uuid": manifest["bundle_uuid"], - "source": manifest["source"], + "source": portable_source_record(manifest["source"], entry), "generation": dict(GENERATION_CONTRACT), "action_token_mapping": manifest["action_token_mapping"], "fast_codec": { @@ -1052,32 +955,46 @@ def effective_fast_config(source_dir: Path) -> dict[str, Any]: raise StarVLAError(f"failed to load FAST config.yaml: {exc}") from exc if not isinstance(source, dict): raise StarVLAError("FAST config.yaml must contain an object") + framework = source.get("framework") + datasets = source.get("datasets") + if not isinstance(framework, dict) or not isinstance(datasets, dict): + raise StarVLAError("FAST config.yaml is missing framework or datasets") + action = framework.get("action_model") + vla = datasets.get("vla_data") + if not isinstance(action, dict) or not isinstance(vla, dict): + raise StarVLAError("FAST config.yaml is missing action_model or vla_data") + action_dim = action.get("action_dim") + future_window = action.get("future_action_window_size") + cot_prompt = vla.get("CoT_prompt") + image_names = vla.get("obs") + image_size = vla.get("image_size") + if type(action_dim) is not int or type(future_window) is not int: + raise StarVLAError("FAST action dimensions must be integers") + if not isinstance(cot_prompt, str) or not cot_prompt: + raise StarVLAError("FAST CoT_prompt must be a non-empty string") + if not isinstance(image_names, list) or not image_names or any( + not isinstance(name, str) or not name for name in image_names + ): + raise StarVLAError("FAST obs must be a non-empty list of image names") + if ( + not isinstance(image_size, list) + or len(image_size) != 2 + or any(type(value) is not int or value <= 0 for value in image_size) + ): + raise StarVLAError("FAST image_size must contain two positive integers") return { "schema_version": 1, - "framework": "QwenFast", + "framework": str(framework.get("framework_py", "QwenFast")), "backbone": BACKBONE, "action_model": "autoregressive_vlm_lm_head", - "action_dim": ACTION_DIM, - "action_horizon": ACTION_HORIZON, - "cot_prompt": COT_PROMPT, - "image_count": 1, - "image_size": [224, 224], + "action_dim": action_dim, + "action_horizon": future_window + 1, + "cot_prompt": cot_prompt, + "image_count": len(image_names), + "image_names": image_names, + "image_size": image_size, "generation": dict(GENERATION_CONTRACT), "source_config_sha256": sha256_file(source_dir / "config.yaml"), - "resolved_overrides": { - "framework.action_model.action_model_type": { - "source": source.get("framework", {}) - .get("action_model", {}) - .get("action_model_type"), - "effective": "FAST", - "authority": "pinned_QwenFast_factory", - }, - "framework.action_model.action_horizon": { - "source": None, - "effective": ACTION_HORIZON, - "authority": "future_action_window_size_plus_current_step", - }, - }, } @@ -1090,12 +1007,10 @@ def stage_checkpoint( staging_dir: Path, catalog: Mapping[str, Any], max_shard_size: int, - verify_hash: bool, ) -> dict[str, Any]: entry, qwen_entry, codec_entry = validate_catalog_contract(catalog) report = preflight(catalog, source_dir, qwen_dir, codec_dir) - if verify_hash: - verify_checkpoint_file(checkpoint, entry) + verify_checkpoint_file(checkpoint, entry) if staging_dir.exists(): raise StarVLAError(f"refusing to overwrite staging directory: {staging_dir}") staging_dir.parent.mkdir(parents=True, exist_ok=True) @@ -1136,20 +1051,18 @@ def stage_checkpoint( codec = validate_fast_codec(codec_dir, codec_entry) manifest = { "schema_version": 1, - "kind": "starvla_qwen25_fast_official_checkpoint_staging", + "kind": "starvla_qwen25_fast_checkpoint_staging", "variant": VARIANT_KEY, "framework": FRAMEWORK, "backbone": BACKBONE, "model_type": MODEL_TYPE, - "bundle_uuid": official_bundle_uuid(entry, catalog), + "bundle_uuid": bundle_uuid(entry, catalog), "source": { "repo_id": entry["repo_id"], "revision": entry["revision"], "checkpoint": str(checkpoint.resolve()), "checkpoint_size": checkpoint.stat().st_size, - "checkpoint_sha256": entry["checkpoint"]["sha256"] - if verify_hash - else sha256_file(checkpoint), + "checkpoint_sha256": entry["checkpoint"]["sha256"], "starvla_revision": catalog["source_revisions"]["starvla"], "llama_cpp_revision": catalog["source_revisions"]["llama_cpp"], "qwen_repo_id": qwen_entry["repo_id"], @@ -1189,12 +1102,12 @@ def validate_staging_manifest( entry, qwen_entry, codec_entry = validate_catalog_contract(catalog) expected = { "schema_version": 1, - "kind": "starvla_qwen25_fast_official_checkpoint_staging", + "kind": "starvla_qwen25_fast_checkpoint_staging", "variant": VARIANT_KEY, "framework": FRAMEWORK, "backbone": BACKBONE, "model_type": MODEL_TYPE, - "bundle_uuid": official_bundle_uuid(entry, catalog), + "bundle_uuid": bundle_uuid(entry, catalog), } mismatches = [ f"{key}: expected {value!r}, got {manifest.get(key)!r}" @@ -1399,6 +1312,7 @@ def convert_staging( ) bundle = build_bundle_manifest( manifest=manifest, + entry=entry, codec=validate_fast_codec(codec_dir, codec_entry), text_component={ "path": TEXT_FILENAME, @@ -1444,7 +1358,6 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: parser.add_argument("--preflight", action="store_true") parser.add_argument("--dry-run", action="store_true") parser.add_argument("--stage-only", action="store_true") - parser.add_argument("--skip-hash-check", action="store_true") return parser.parse_args(argv) @@ -1462,7 +1375,7 @@ def main(argv: Sequence[str] | None = None) -> int: print(json.dumps(report, indent=2, sort_keys=True)) return 0 if args.dry_run: - output_dir = args.output_dir or Path("ckpts/starvla/gguf/qwen25-fast") + output_dir = args.output_dir or Path("ckpts/starvla/gguf/qwen25_fast") commands = build_commands( args.python, args.staging_dir / "hf", @@ -1506,7 +1419,6 @@ def main(argv: Sequence[str] | None = None) -> int: staging_dir=args.staging_dir, catalog=catalog, max_shard_size=args.max_shard_size, - verify_hash=not args.skip_hash_check, ) print(f"staging manifest: {args.staging_dir / STAGING_MANIFEST_FILENAME}") if args.stage_only: diff --git a/tools/hf2gguf/starvla/convert_starvla_qwen_to_gguf.py b/tools/hf2gguf/starvla/convert_starvla_qwen_to_gguf.py index 8a24289..4243dcb 100755 --- a/tools/hf2gguf/starvla/convert_starvla_qwen_to_gguf.py +++ b/tools/hf2gguf/starvla/convert_starvla_qwen_to_gguf.py @@ -20,9 +20,8 @@ default_text_filename, get_variant, load_catalog, - validate_official_surgery_manifest, + validate_surgery_manifest, verify_staged_assets, - verify_staged_tensors_against_checkpoint, ) @@ -116,13 +115,13 @@ def verify_llama_checkout(path: Path, expected_revision: str) -> Path: if actual_revision != expected_revision: raise StarVLAError( f"llama.cpp revision mismatch: expected {expected_revision}, got {actual_revision}; " - "update the pinned catalog and regenerate golden data before converting" + "use the revision pinned by the checkpoint catalog" ) worktree_changes = git_worktree_changes(root) if worktree_changes: raise StarVLAError( "llama.cpp has tracked or untracked worktree changes; " - f"use the clean pinned revision for official conversion:\n{worktree_changes}" + f"use the clean pinned revision for conversion:\n{worktree_changes}" ) return root @@ -172,7 +171,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--surgery-manifest", type=Path, required=True) parser.add_argument("--output-dir", type=Path, required=True) parser.add_argument("--catalog", type=Path, default=DEFAULT_CATALOG) - parser.add_argument("--llama-root", type=Path, default=LLAMA_ROOT) + parser.add_argument("--llama-root", type=Path, required=True) parser.add_argument("--text-filename") parser.add_argument("--mmproj-filename") parser.add_argument( @@ -203,16 +202,8 @@ def main() -> int: catalog = load_catalog(args.catalog) variant_name = str(manifest.get("variant", "")) variant = get_variant(catalog, variant_name) - validate_official_surgery_manifest(manifest, variant, catalog) + validate_surgery_manifest(manifest, variant, catalog) verify_staged_assets(args.hf_dir, manifest.get("qwen_assets", {}), component="Qwen") - verify_staged_tensors_against_checkpoint( - args.hf_dir, - manifest.get("vlm_output", {}), - manifest, - variant, - component="vlm", - ) - expected_revision = str(manifest.get("source", {}).get("llama_cpp_revision", "")) llama_root = verify_llama_checkout(args.llama_root, expected_revision) @@ -227,7 +218,7 @@ def main() -> int: mmproj_metadata = args.output_dir / "mmproj-metadata.json" bundle_uuid = str(manifest["bundle_uuid"]) source = manifest["source"] - backbone = str(manifest.get("backbone", variant.get("backbone", "qwen3_vl"))) + backbone = str(manifest.get("backbone", variant["backbone"])) backbone_label = { "qwen3_vl": "Qwen3-VL", "qwen2_5_vl": "Qwen2.5-VL", diff --git a/tools/hf2gguf/starvla/download_starvla.py b/tools/hf2gguf/starvla/download_starvla.py index ceb712f..364f2b5 100755 --- a/tools/hf2gguf/starvla/download_starvla.py +++ b/tools/hf2gguf/starvla/download_starvla.py @@ -10,8 +10,6 @@ from starvla_checkpoint import ( DEFAULT_CATALOG, - DEFAULT_QWEN_ASSET, - SUPPORTED_BACKBONES, StarVLAError, atomic_write_json, get_variant, @@ -21,72 +19,6 @@ DEFAULT_BACKBONE = "qwen3_vl" -DEFAULT_TARGET_MATRIX = Path(__file__).with_name("release_targets.json") - - -def load_target_matrix(path: Path | str = DEFAULT_TARGET_MATRIX) -> dict[str, Any]: - import json - - matrix_path = Path(path) - try: - matrix = json.loads(matrix_path.read_text(encoding="utf-8")) - except (OSError, json.JSONDecodeError) as exc: - raise StarVLAError(f"failed to load release target matrix {matrix_path}: {exc}") from exc - if not isinstance(matrix, dict) or matrix.get("schema_version") != 1: - raise StarVLAError("unsupported StarVLA release target matrix") - return matrix - - -def validate_target_matrix( - catalog: dict[str, Any], matrix: dict[str, Any] -) -> None: - variants = catalog["variants"] - expected_backbones = set(SUPPORTED_BACKBONES) - seen: set[str] = set() - for tier in ("targets", "experimental"): - groups = matrix.get(tier) - if not isinstance(groups, dict) or set(groups) != expected_backbones: - raise StarVLAError( - f"release target matrix {tier} must cover every supported backbone exactly" - ) - for backbone, names in groups.items(): - if ( - not isinstance(names, list) - or any(not isinstance(name, str) for name in names) - or len(names) != len(set(names)) - ): - raise StarVLAError( - f"release target matrix {tier}.{backbone} must be a unique list" - ) - if tier == "targets" and not names: - raise StarVLAError( - f"release target matrix targets.{backbone} cannot be empty" - ) - for name in names: - entry = variants.get(name) - if not isinstance(entry, dict): - raise StarVLAError(f"release target {name!r} is not a catalog variant") - if variant_backbone(entry) != backbone: - raise StarVLAError(f"release target {name!r} has the wrong backbone") - if name in seen: - raise StarVLAError(f"release target {name!r} occurs more than once") - if ( - entry.get("status") != "official_policy" - or entry.get("checkpoint") is None - ): - raise StarVLAError( - f"release target {name!r} is not an official policy checkpoint" - ) - seen.add(name) - policy_variants = { - name - for name, entry in variants.items() - if entry.get("status") == "official_policy" and entry.get("checkpoint") is not None - } - if seen != policy_variants: - raise StarVLAError( - "release target matrix must classify every policy variant exactly once" - ) def destination_for(root: Path, entry: dict[str, Any]) -> Path: @@ -94,7 +26,7 @@ def destination_for(root: Path, entry: dict[str, Any]) -> Path: def variant_backbone(entry: dict[str, Any]) -> str: - return str(entry.get("backbone", DEFAULT_BACKBONE)) + return str(entry["backbone"]) def available_backbones(catalog: dict[str, Any]) -> list[str]: @@ -110,7 +42,6 @@ def resolve_variant_keys( catalog: dict[str, Any], requested: Sequence[str] | None, requested_backbone: str | None, - target_matrix: dict[str, Any] | None = None, ) -> tuple[str, list[str]]: variants = catalog["variants"] backbones = available_backbones(catalog) @@ -138,15 +69,11 @@ def resolve_variant_keys( if variant_backbone(entry) == backbone } tokens = list(requested or ("oft",)) - if "all" in tokens or "catalog-all" in tokens: + if "all" in tokens: if len(tokens) != 1: raise StarVLAError( - "--variant all/catalog-all cannot be combined with another variant" + "--variant all cannot be combined with another variant" ) - if tokens[0] == "all": - matrix = target_matrix or load_target_matrix() - validate_target_matrix(catalog, matrix) - return backbone, list(matrix["targets"][backbone]) return backbone, list(candidates) selected: list[str] = [] @@ -165,7 +92,6 @@ def resolve_variant_keys( *candidates, *(str(entry["framework"]) for entry in candidates.values()), "all", - "catalog-all", } ) raise StarVLAError( @@ -191,7 +117,7 @@ def required_shared_assets( has_fast = False for variant_key in variant_keys: entry = get_variant(catalog, variant_key) - qwen_asset = str(entry.get("qwen_asset", DEFAULT_QWEN_ASSET)) + qwen_asset = str(entry["qwen_asset"]) if qwen_asset not in catalog["shared_assets"]: raise StarVLAError( f"variant {variant_key!r} references unknown Qwen asset {qwen_asset!r}" @@ -217,7 +143,7 @@ def download_entry( result = { "repo_id": entry["repo_id"], "revision": entry["revision"], - "directory": str(destination), + "directory": (Path(str(entry["directory"])) / str(entry["revision"])).as_posix(), "requested_files": files, "files": [], } @@ -274,8 +200,8 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: action="append", help=( "catalog variant key or framework alias to download " - "(repeatable; 'all' selects the release matrix, 'catalog-all' also " - "includes experimental entries; default: OFT for the selected backbone)" + "(repeatable; 'all' selects every variant for the backbone; " + "default: OFT for the selected backbone)" ), ) parser.add_argument( @@ -287,15 +213,16 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: ) parser.add_argument("--root", type=Path, default=Path("ckpts/starvla/sources")) parser.add_argument("--catalog", type=Path, default=DEFAULT_CATALOG) - parser.add_argument( - "--target-matrix", type=Path, default=DEFAULT_TARGET_MATRIX, - help="release/experimental support matrix used by --variant all", - ) parser.add_argument( "--metadata-only", action="store_true", help="skip policy checkpoints and all optional base weights", ) + parser.add_argument( + "--skip-checkpoint", + action="store_true", + help="skip policy checkpoints while retaining requested optional base weights", + ) parser.add_argument( "--include-base-weights", action="store_true", @@ -305,8 +232,8 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: "--include-fast-weights", action="store_true", help=( - "download the action-ready base weights for a selected FAST variant " - "(kept for compatibility; the policy checkpoint is downloaded separately)" + "download the action-ready Qwen weights for a selected FAST variant; " + "the policy checkpoint is downloaded separately" ), ) parser.add_argument("--no-shared-assets", action="store_true") @@ -320,21 +247,16 @@ def main() -> int: args = parse_args() try: catalog = load_catalog(args.catalog) - target_matrix = load_target_matrix(args.target_matrix) - validate_target_matrix(catalog, target_matrix) - backbone, variants = resolve_variant_keys( - catalog, args.variant, args.backbone, target_matrix - ) + backbone, variants = resolve_variant_keys(catalog, args.variant, args.backbone) manifest: dict[str, Any] = { "schema_version": 1, - "catalog": str(args.catalog.resolve()), - "target_matrix": str(args.target_matrix.resolve()), - "target_matrix_sha256": sha256_file(args.target_matrix), + "catalog_sha256": sha256_file(args.catalog), "source_revisions": catalog["source_revisions"], "backbone": backbone, "variants": variants, "metadata_only": bool(args.metadata_only), + "skip_checkpoint": bool(args.skip_checkpoint), "downloads": {}, } @@ -345,7 +267,7 @@ def main() -> int: for variant in variants } fast_qwen_assets = { - str(entry.get("qwen_asset", DEFAULT_QWEN_ASSET)) + str(entry["qwen_asset"]) for entry in variant_entries.values() if entry.get("framework") == "fast" } @@ -371,15 +293,9 @@ def main() -> int: for variant in variants: entry = get_variant(catalog, variant) files = list(entry.get("files", [])) - checkpoint = entry.get("checkpoint") - if checkpoint is not None and not args.metadata_only: + checkpoint = entry["checkpoint"] + if not args.metadata_only and not args.skip_checkpoint: files.append(str(checkpoint["path"])) - include_variant_weights = ( - args.include_base_weights - or (entry.get("framework") == "fast" and args.include_fast_weights) - ) - if include_variant_weights and not args.metadata_only: - files.extend(entry.get("optional_weight_files", [])) download = download_entry( entry, args.root, @@ -390,7 +306,7 @@ def main() -> int: ) manifest["downloads"][f"variant:{variant}"] = download - if checkpoint is not None and not args.metadata_only and not args.dry_run: + if not args.metadata_only and not args.skip_checkpoint and not args.dry_run: record = next( (item for item in download["files"] if item["path"] == checkpoint["path"]), None, diff --git a/tools/hf2gguf/starvla/environment.yaml b/tools/hf2gguf/starvla/environment.yaml index 2deb0dd..5052ab7 100644 --- a/tools/hf2gguf/starvla/environment.yaml +++ b/tools/hf2gguf/starvla/environment.yaml @@ -5,10 +5,19 @@ dependencies: - python=3.11 - pip - pip: - - torch - - numpy - - safetensors + - torch==2.6.0 + - torchvision==0.21.0 + - numpy==1.26.4 + - safetensors==0.7.0 - sentencepiece - transformers==4.57.0 + - tokenizers==0.22.2 + - accelerate==1.5.2 + - diffusers==0.37.1 + - omegaconf==2.3.0 + - pillow==12.1.1 + - qwen-vl-utils==0.0.14 + - rich + - scipy - huggingface_hub>=0.36.0 - pyyaml diff --git a/tools/hf2gguf/starvla/generate_starvla_groot_golden.py b/tools/hf2gguf/starvla/generate_starvla_groot_golden.py deleted file mode 100644 index d90ede7..0000000 --- a/tools/hf2gguf/starvla/generate_starvla_groot_golden.py +++ /dev/null @@ -1,1169 +0,0 @@ -#!/usr/bin/env python3 -"""Generate an independent, fixed-noise oracle for official Qwen-GR00T.""" - -from __future__ import annotations - -import argparse -import datetime as dt -import gc -import hashlib -import importlib.metadata -import json -import os -import platform -import random -import shutil -import sys -import tempfile -from pathlib import Path -from typing import Any, Iterable, Mapping, Sequence - -import numpy as np - - -TOOLS_DIR = Path(__file__).resolve().parent -if str(TOOLS_DIR) not in sys.path: - sys.path.insert(0, str(TOOLS_DIR)) - -from generate_starvla_pi_v3_golden import ( # noqa: E402 - EXPECTED_ACCELERATE_VERSION, - EXPECTED_DIFFUSERS_VERSION, - EXPECTED_NUMPY_VERSION, - EXPECTED_OMEGACONF_VERSION, - EXPECTED_PILLOW_VERSION, - EXPECTED_QWEN2VL_IMAGE_PROCESSING_FAST_SHA256, - EXPECTED_QWEN2VL_IMAGE_PROCESSING_SHA256, - EXPECTED_QWEN3VL_MODELING_SHA256, - EXPECTED_QWEN3VL_PROCESSING_SHA256, - EXPECTED_SAFETENSORS_VERSION, - EXPECTED_TOKENIZERS_VERSION, - EXPECTED_TORCHVISION_VERSION, - EXPECTED_TORCH_VERSION, - EXPECTED_TRANSFORMERS_GENERIC_SHA256, - EXPECTED_TRANSFORMERS_VERSION, - OFFICIAL_ENVIRONMENT_FREEZE, - _array_record, - _array_sha256, - _assert_module_origin, - _canonical_json, - _config_only_qwen_bootstrap, - _configure_determinism, - _official_qwen_model_alias, - _sha256_bytes, - validate_runtime_versions, - verify_pinned_source_checkout, - verify_transformers_qwen3vl_recorder_semantics, -) -from starvla_checkpoint import ( # noqa: E402 - DEFAULT_CATALOG, - StarVLAError, - get_variant, - load_catalog, - official_bundle_uuid, - resolve_effective_config, - sha256_file, - verify_catalog_files, - verify_checkpoint_file, -) - - -GOLDEN_SCHEMA_VERSION = 1 -GOLDEN_KIND = "starvla_groot_official_python_oracle" -RUNNER_CONTRACT_KIND = "starvla_groot_runner_contract" -SUPPORTED_VARIANT = "groot" -SEED = 0 -EXPECTED_ACTION_HORIZON = 16 -EXPECTED_ACTION_DIM = 7 -EXPECTED_QWEN_HIDDEN_DIM = 2560 -EXPECTED_DIT_WIDTH = 768 -EXPECTED_DIT_OUTPUT_DIM = 1024 -EXPECTED_DIT_BLOCK_COUNT = 16 -EXPECTED_FUTURE_TOKEN_COUNT = 32 -EXPECTED_SEQUENCE_LENGTH = 48 -EXPECTED_TIMESTEP_IDS = [0, 250, 500, 750] -EXPECTED_COT_TEMPLATE = ( - "Your task is {instruction}. To identify the key objects for your task. " - "Locate their bounding boxes in [x1,y1,x2,y2] format." -) - -PINNED_SOURCE_FILES = { - "starVLA/model/framework/VLM4A/QwenGR00T.py": - "645d99d8d6a8daaccb7bb6e3211971b5cc7396d39968b0e9c20c3894d6883249", - "starVLA/model/modules/action_model/GR00T_ActionHeader.py": - "a01c7ca048589835a23bf46cf670275dfa643a1fb2da0bafd14654e1a57236e5", - "starVLA/model/modules/action_model/flow_matching_head/cross_attention_dit.py": - "c18d2e128dddcd67dc88c4fb178c99d7ceb7ea40d40ea9622b120151b81db359", - "starVLA/model/modules/vlm/QWen3.py": - "03e0c35cfe86490886ff26a59230f27726ba4b46259d2be20beed4c532925d47", - "deployment/model_server/policy_norm_processor.py": - "3fd280c8f5072943fad6809dd5705cb713007c10d7240d2a23e3dadcd3963d2a", -} - -ARRAY_SOURCE_DTYPES = { - "input_ids": "int64", - "attention_mask": "bool", - "image_grid_thw": "int64", - "raw_l_out_35": "bfloat16", - "initial_noise": "bfloat16", - "action_features": "float32", - "dit_inputs": "float32", - "dit_block_outputs": "float32", - "dit_outputs": "float32", - "predicted_velocities": "float32", - "actions_after_steps": "float32", - "normalized_actions": "float32", - "unnormalized_actions": "float32", -} - - -def _distribution_version(name: str) -> str: - try: - return importlib.metadata.version(name) - except importlib.metadata.PackageNotFoundError as exc: - raise StarVLAError(f"required package is not installed: {name}") from exc - - -def _regular_file(path: Path, *, label: str) -> Path: - if path.is_symlink() or not path.is_file(): - raise StarVLAError(f"{label} must be a regular, non-symlink file: {path}") - return path.resolve() - - -def _source_asset_hashes(entry: Mapping[str, Any], *, staged: bool = False) -> dict[str, str]: - overrides = entry.get("staged_overrides", {}) if staged else {} - return { - relative: overrides.get(relative, record)["sha256"] - for relative, record in entry["file_hashes"].items() - } - - -def verify_source_semantics(source_dir: Path) -> dict[str, Any]: - """Bind the sampler schedule and ordering to the pinned official source.""" - - actual: dict[str, str] = {} - for relative, expected_hash in PINNED_SOURCE_FILES.items(): - path = _regular_file(source_dir / relative, label=f"pinned source {relative}") - digest = sha256_file(path) - if digest != expected_hash: - raise StarVLAError( - f"pinned GR00T source SHA256 mismatch for {relative}: " - f"expected {expected_hash}, got {digest}" - ) - actual[relative] = digest - - action_source = (source_dir / "starVLA/model/modules/action_model/GR00T_ActionHeader.py").read_text( - encoding="utf-8" - ) - qwen_source = (source_dir / "starVLA/model/framework/VLM4A/QwenGR00T.py").read_text( - encoding="utf-8" - ) - required_action_fragments = ( - "dtype=vl_embs.dtype", - "dt = 1.0 / num_steps", - "t_cont = t / float(num_steps)", - "t_discretized = int(t_cont * self.num_timestep_buckets)", - "torch.cat((future_tokens, action_features), dim=1)", - "pred_velocity = pred[:, -self.action_horizon :]", - "actions = actions + dt * pred_velocity", - ) - required_qwen_fragments = ( - "backbone_attention_mask = backbone_attention_mask.to(dtype=torch.bool)", - "last_hidden = qwenvl_outputs.hidden_states[-1]", - "self.action_model.predict_action(", - "last_hidden, state, encoder_attention_mask=backbone_attention_mask", - ) - missing = [value for value in required_action_fragments if value not in action_source] - missing += [value for value in required_qwen_fragments if value not in qwen_source] - if missing: - raise StarVLAError(f"pinned GR00T source semantics probe failed: missing {missing!r}") - return { - "files": actual, - "schedule_source": "GR00T_ActionHeader.FlowmatchingActionHead.predict_action", - "continuous_formula": "t / float(num_steps)", - "bucket_formula": "int(t_cont * self.num_timestep_buckets)", - "observed_expected_timestep_ids": EXPECTED_TIMESTEP_IDS, - "query_sequence": "future_tokens_then_action_features", - "conditioning_sequence": "complete_qwen_outer_raw_l_out_35_with_bool_attention_mask", - "velocity_slice": "last_action_horizon_tokens", - "euler_update": "actions = actions + dt * pred_velocity", - } - - -def validate_available_inputs( - *, checkpoint_root: Path, source_dir: Path, catalog_path: Path = DEFAULT_CATALOG -) -> dict[str, Any]: - catalog = load_catalog(catalog_path) - variant = get_variant(catalog, SUPPORTED_VARIANT) - qwen = catalog["shared_assets"]["qwen3_vl_4b_instruct"] - checkpoint_root = checkpoint_root.resolve() - expected_source = (checkpoint_root / "source" / "starvla").resolve() - if source_dir.resolve() != expected_source: - raise StarVLAError( - f"StarVLA source must be the canonical checkout {expected_source}, got {source_dir.resolve()}" - ) - verify_pinned_source_checkout(source_dir, catalog["source_revisions"]["starvla"]) - source_probe = verify_source_semantics(source_dir) - policy_dir = checkpoint_root / "sources" / variant["directory"] / variant["revision"] - qwen_dir = checkpoint_root / "sources" / qwen["directory"] / qwen["revision"] - verify_catalog_files(policy_dir, variant) - verify_catalog_files(qwen_dir, qwen) - checkpoint = policy_dir / variant["checkpoint"]["path"] - sidecar = Path(f"{checkpoint}.aria2") - checkpoint_ready = checkpoint.is_file() and not checkpoint.is_symlink() and not sidecar.exists() - if checkpoint_ready: - verify_checkpoint_file(checkpoint, variant) - return { - "catalog": catalog, - "catalog_path": catalog_path.resolve(), - "variant": variant, - "qwen": qwen, - "policy_dir": policy_dir, - "qwen_dir": qwen_dir, - "checkpoint": checkpoint, - "checkpoint_ready": checkpoint_ready, - "source_dir": source_dir.resolve(), - "source_probe": source_probe, - } - - -def _validate_effective_config(config: Mapping[str, Any]) -> None: - try: - framework = config["framework"] - action = framework["action_model"] - diffusion = action["diffusion_model_cfg"] - vla = config["datasets"]["vla_data"] - except (KeyError, TypeError) as exc: - raise StarVLAError("effective GR00T config is missing required objects") from exc - actual = { - "framework": framework.get("name"), - "action_model_type": action.get("action_model_type"), - "action_horizon": action.get("action_horizon"), - "action_dim": action.get("action_dim"), - "state_dim": action.get("state_dim"), - "steps": action.get("num_inference_timesteps"), - "buckets": action.get("num_timestep_buckets"), - "future_tokens": action.get("num_target_vision_tokens"), - "width": diffusion.get("input_embedding_dim"), - "layers": diffusion.get("num_layers"), - "heads": diffusion.get("num_attention_heads"), - "head_dim": diffusion.get("attention_head_dim"), - "cross_dim": diffusion.get("cross_attention_dim"), - "output_dim": diffusion.get("output_dim"), - "interleave": diffusion.get("interleave_self_attention"), - "image_size": vla.get("image_size"), - "obs_image_size": vla.get("obs_image_size"), - "obs": vla.get("obs"), - "data_mix": vla.get("data_mix"), - "cot": vla.get("CoT_prompt"), - } - expected = { - "framework": "QwenGR00T", - "action_model_type": "DiT-B", - "action_horizon": 16, - "action_dim": 7, - "state_dim": 7, - "steps": 4, - "buckets": 1000, - "future_tokens": 32, - "width": 768, - "layers": 16, - "heads": 12, - "head_dim": 64, - "cross_dim": 2560, - "output_dim": 1024, - "interleave": True, - "image_size": [224, 224], - "obs_image_size": None, - "obs": ["image_0"], - "data_mix": "bridge_rt_1", - "cot": EXPECTED_COT_TEMPLATE, - } - if actual != expected: - raise StarVLAError(f"unexpected effective official GR00T config: {actual}") - - -def expected_model_instruction(config: Mapping[str, Any], task: str) -> str: - if not isinstance(task, str) or not task or "\x00" in task: - raise StarVLAError("task must be a non-empty string without NUL") - template = config["datasets"]["vla_data"]["CoT_prompt"] - if template != EXPECTED_COT_TEMPLATE or template.count("{instruction}") != 1: - raise StarVLAError("official GR00T CoT prompt contract changed") - return template.replace("{instruction}", task) - - -def load_official_framework(paths: Mapping[str, Any], *, device: str) -> tuple[Any, dict[str, Any]]: - import torch - import transformers - - if not paths["checkpoint_ready"]: - raise StarVLAError( - f"official GR00T checkpoint is absent or incomplete: {paths['checkpoint']}" - ) - source_dir = Path(paths["source_dir"]) - if any(name == "starVLA" or name.startswith("starVLA.") for name in sys.modules): - raise StarVLAError("starVLA was imported before pinned-source verification") - sys.path.insert(0, str(source_dir)) - try: - from starVLA.model.framework import base_framework, share_tools - from starVLA.model.framework.VLM4A import QwenGR00T - - for module in (base_framework, share_tools, QwenGR00T): - _assert_module_origin(module, source_dir) - config = resolve_effective_config(Path(paths["policy_dir"]), SUPPORTED_VARIANT) - _validate_effective_config(config) - qwen_dir = Path(paths["qwen_dir"]).resolve() - with _official_qwen_model_alias(qwen_dir) as qwen_alias: - config = base_framework.merge_config_overrides( - config, - [ - f"framework.qwenvl.base_vlm={qwen_alias}", - "framework.qwenvl.attn_implementation=sdpa", - ], - ) - configured_qwen = Path(config["framework"]["qwenvl"]["base_vlm"]) - if "Qwen3-VL" not in str(configured_qwen) or configured_qwen.resolve() != qwen_dir: - raise StarVLAError( - "effective GR00T Qwen source does not preserve official dispatch and pinned assets" - ) - cfg = share_tools.dict_to_namespace(config) - cfg.trainer.pretrained_checkpoint = None - with _config_only_qwen_bootstrap(torch, transformers, qwen_dir): - framework = QwenGR00T.Qwen_GR00T(cfg) - try: - state = torch.load(paths["checkpoint"], map_location="cpu", mmap=True, weights_only=True) - except TypeError: - state = torch.load(paths["checkpoint"], map_location="cpu", weights_only=True) - if not isinstance(state, Mapping) or not state: - raise StarVLAError("official GR00T checkpoint did not contain a state_dict") - framework.load_state_dict(state, strict=True) - del state - gc.collect() - action_model = framework.action_model - if type(framework).__name__ != "Qwen_GR00T": - raise StarVLAError(f"unexpected official framework class: {type(framework).__name__}") - if len(action_model.model.transformer_blocks) != EXPECTED_DIT_BLOCK_COUNT: - raise StarVLAError("official GR00T DiT block count is not 16") - if int(framework.action_horizon) != EXPECTED_ACTION_HORIZON: - raise StarVLAError("official GR00T action horizon changed") - qwen_dtypes = {parameter.dtype for parameter in framework.qwen_vl_interface.parameters()} - policy_dtypes = {parameter.dtype for parameter in action_model.parameters()} - if qwen_dtypes != {torch.bfloat16} or policy_dtypes != {torch.float32}: - raise StarVLAError( - "official GR00T strict-load dtype boundary changed: " - f"qwen={qwen_dtypes}, policy={policy_dtypes}" - ) - return framework.to(device).eval(), config - finally: - if sys.path and sys.path[0] == str(source_dir): - del sys.path[0] - - -def _tensor_to_array(tensor: Any) -> tuple[np.ndarray, str]: - source_dtype = str(tensor.dtype).removeprefix("torch.") - value = tensor.detach().cpu().contiguous() - if source_dtype == "bfloat16": - value = value.float() - return np.ascontiguousarray(value.numpy()), source_dtype - - -def run_official_forward( - framework: Any, *, images: Sequence[Any], task: str, seed: int = SEED -) -> dict[str, Any]: - """Run pinned source with only the declared BF16-to-FP32 compatibility widens.""" - - import torch - - captures: dict[str, Any] = {} - qwen = framework.qwen_vl_interface - action_model = framework.action_model - language_model = qwen.model.model.language_model - raw_final: dict[str, Any] = {} - raw_qwen_taps: list[Any | None] = [None] * len(language_model.layers) - language_inputs: dict[str, Any] = {} - result_norm: dict[str, Any] = {} - block_outputs: list[Any] = [] - handles = [] - original_build = qwen.build_qwenvl_inputs - original_policy = action_model.predict_action - original_action_encoder = action_model.action_encoder.forward - original_dit = action_model.model.forward - original_decoder = action_model.action_decoder.forward - - def capture_build(*args: Any, **kwargs: Any): - if "qwen_inputs" in captures: - raise StarVLAError("official GR00T preprocessing ran more than once") - batch_images = kwargs.get("images", args[0] if args else None) - instructions = kwargs.get("instructions", args[1] if len(args) > 1 else None) - captures["processed_images"] = list(batch_images[0]) - captures["framework_instructions"] = list(instructions) - output = original_build(*args, **kwargs) - captures["qwen_inputs"] = { - key: value.detach() for key, value in output.items() if isinstance(value, torch.Tensor) - } - return output - - def capture_outer(_module: Any, _inputs: Any, output: Any): - hidden = getattr(output, "hidden_states", None) - if hidden is None or len(hidden) != 37: - raise StarVLAError("official Qwen outer recorder did not expose 37 hidden tuple entries") - captures["outer_raw_final"] = hidden[-1].detach().clone() - - def capture_language_inputs(_module: Any, args: Any, kwargs: Any): - if language_inputs: - raise StarVLAError("official Qwen language model ran more than once") - if args: - raise StarVLAError("official Qwen language model stopped using keyword inputs") - inputs_embeds = kwargs.get("inputs_embeds") - visual_pos_masks = kwargs.get("visual_pos_masks") - deepstack_visual_embeds = kwargs.get("deepstack_visual_embeds") - if (inputs_embeds is None or visual_pos_masks is None or - deepstack_visual_embeds is None): - raise StarVLAError("official Qwen language model omitted prepared visual inputs") - language_inputs["inputs_embeds"] = inputs_embeds.detach().clone() - language_inputs["visual_pos_masks"] = visual_pos_masks.detach().clone() - language_inputs["deepstack_visual_embeds"] = [ - value.detach().clone() for value in deepstack_visual_embeds - ] - - def capture_policy(*args: Any, **kwargs: Any): - vl_embs = args[0] if args else kwargs.get("vl_embs") - state = args[1] if len(args) > 1 else kwargs.get("state") - policy_mask = kwargs.get( - "encoder_attention_mask", args[2] if len(args) > 2 else None - ) - if state is not None: - raise StarVLAError("official GR00T oracle unexpectedly entered the state branch") - if policy_mask is None: - raise StarVLAError("official GR00T policy did not receive an attention mask") - captures["policy_qwen_input"] = vl_embs.detach().clone() - captures["policy_attention_mask"] = policy_mask.detach().clone() - original_randn = torch.randn - - def capture_randn(*randn_args: Any, **randn_kwargs: Any): - value = original_randn(*randn_args, **randn_kwargs) - if "initial_noise" in captures: - raise StarVLAError("official GR00T policy sampled initial noise more than once") - captures["initial_noise"] = value.detach().clone() - return value - - torch.randn = capture_randn - try: - output = original_policy(*args, **kwargs) - finally: - torch.randn = original_randn - captures["raw_policy"] = output.detach().clone() - return output - - def capture_action_encoder(actions: Any, timesteps: Any): - captures.setdefault("action_inputs", []).append(actions.detach().clone()) - output = original_action_encoder(actions.to(dtype=torch.float32), timesteps) - captures.setdefault("action_encoder_outputs", []).append(output.detach().clone()) - return output - - def capture_dit(*args: Any, **kwargs: Any): - hidden = kwargs.get("hidden_states", args[0] if args else None) - conditioning = kwargs.get("encoder_hidden_states", args[1] if len(args) > 1 else None) - timestep = kwargs.get("timestep", args[2] if len(args) > 2 else None) - captures.setdefault("dit_inputs", []).append(hidden.detach().clone()) - captures.setdefault("dit_conditioning_inputs", []).append(conditioning.detach().clone()) - captures.setdefault("timestep_ids", []).append(int(timestep.item())) - if "encoder_hidden_states" in kwargs: - kwargs["encoder_hidden_states"] = conditioning.to(dtype=torch.float32) - else: - args = list(args) - args[1] = conditioning.to(dtype=torch.float32) - args = tuple(args) - output = original_dit(*args, **kwargs) - captures.setdefault("dit_outputs", []).append(output.detach().clone()) - return output - - def capture_decoder(value: Any): - captures.setdefault("decoder_inputs", []).append(value.detach().clone()) - output = original_decoder(value) - captures.setdefault("decoder_outputs", []).append(output.detach().clone()) - return output - - def capture_raw_qwen_tap(layer_index: int): - def capture(_module: Any, _inputs: Any, output: Any): - value = output.detach().clone() - raw_qwen_taps[layer_index] = value - if layer_index + 1 == len(raw_qwen_taps): - raw_final["value"] = value - - return capture - - for layer_index, layer in enumerate(language_model.layers): - handles.append(layer.register_forward_hook(capture_raw_qwen_tap(layer_index))) - handles.append( - language_model.register_forward_pre_hook( - capture_language_inputs, with_kwargs=True - ) - ) - handles.append( - language_model.norm.register_forward_hook( - lambda _m, _i, output: result_norm.__setitem__("value", output.detach().clone()) - ) - ) - handles.append(qwen.model.register_forward_hook(capture_outer)) - for block in action_model.model.transformer_blocks: - handles.append( - block.register_forward_hook( - lambda _m, _i, output: block_outputs.append(output.detach().clone()) - ) - ) - qwen.build_qwenvl_inputs = capture_build - action_model.predict_action = capture_policy - action_model.action_encoder.forward = capture_action_encoder - action_model.model.forward = capture_dit - action_model.action_decoder.forward = capture_decoder - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - try: - result = framework.predict_action(examples=[{"image": list(images), "lang": task}]) - finally: - for handle in handles: - handle.remove() - qwen.build_qwenvl_inputs = original_build - action_model.predict_action = original_policy - action_model.action_encoder.forward = original_action_encoder - action_model.model.forward = original_dit - action_model.action_decoder.forward = original_decoder - - required = { - "qwen_inputs", "processed_images", "framework_instructions", "outer_raw_final", - "policy_qwen_input", "policy_attention_mask", "initial_noise", "raw_policy", "action_inputs", - "action_encoder_outputs", "dit_inputs", "dit_conditioning_inputs", "timestep_ids", - "dit_outputs", "decoder_inputs", "decoder_outputs", - } - missing = sorted(required - set(captures)) - if (missing or "value" not in raw_final or "value" not in result_norm or - any(value is None for value in raw_qwen_taps) or not language_inputs): - raise StarVLAError(f"official GR00T instrumentation missed captures: {missing}") - if captures["framework_instructions"] != [task]: - raise StarVLAError("official GR00T framework instruction changed before Qwen preprocessing") - if len(captures["processed_images"]) != len(images) or any( - actual.mode != expected.mode or actual.size != expected.size or - actual.tobytes() != expected.tobytes() - for actual, expected in zip(captures["processed_images"], images) - ): - raise StarVLAError("official GR00T unexpectedly pre-resized or altered the input image") - if captures["timestep_ids"] != EXPECTED_TIMESTEP_IDS: - raise StarVLAError( - f"official source-derived GR00T timestep order changed: {captures['timestep_ids']}" - ) - if len(block_outputs) != 4 * EXPECTED_DIT_BLOCK_COUNT: - raise StarVLAError("official GR00T instrumentation missed DiT block outputs") - for name in ( - "action_inputs", "action_encoder_outputs", "dit_inputs", "dit_conditioning_inputs", - "dit_outputs", "decoder_inputs", "decoder_outputs", - ): - if len(captures[name]) != 4: - raise StarVLAError(f"official GR00T {name} did not run exactly four times") - if not torch.equal(captures["outer_raw_final"], raw_final["value"]): - raise StarVLAError("outer hidden_states[-1] is not complete raw l_out-35") - if torch.equal(captures["outer_raw_final"], result_norm["value"]): - raise StarVLAError("GR00T conditioning unexpectedly uses result_norm") - if not torch.equal(captures["policy_qwen_input"], captures["outer_raw_final"]): - raise StarVLAError("GR00T policy did not receive complete raw l_out-35") - mask = captures["qwen_inputs"].get("attention_mask") - if mask is None or mask.dtype not in (torch.int64, torch.bool): - raise StarVLAError("official Qwen attention mask source dtype changed") - policy_mask = captures["policy_attention_mask"] - if policy_mask.dtype != torch.bool or not torch.equal(policy_mask, mask.to(dtype=torch.bool)): - raise StarVLAError("official GR00T policy mask is not the complete Qwen boolean mask") - if captures["initial_noise"].dtype != torch.bfloat16: - raise StarVLAError("official GR00T initial noise is no longer sampled as BF16") - if not torch.equal(captures["initial_noise"], captures["action_inputs"][0]): - raise StarVLAError("first GR00T action encoder input is not initial noise") - if captures["outer_raw_final"].dtype != torch.bfloat16: - raise StarVLAError("official raw l_out-35 boundary is no longer BF16") - if any(value.dtype != torch.bfloat16 for value in raw_qwen_taps): - raise StarVLAError("official raw Qwen layer taps are no longer BF16") - if not torch.equal(raw_qwen_taps[-1], captures["outer_raw_final"]): - raise StarVLAError("official raw Qwen layer taps do not end at raw l_out-35") - if any(value.dtype != torch.float32 for value in captures["action_encoder_outputs"]): - raise StarVLAError("GR00T action encoder compatibility output is not FP32") - for name in ("dit_inputs", "dit_outputs", "decoder_inputs", "decoder_outputs"): - if any(value.dtype != torch.float32 for value in captures[name]): - raise StarVLAError(f"GR00T compatibility path {name} is not FP32") - if any(value.dtype != torch.bfloat16 for value in captures["dit_conditioning_inputs"]): - raise StarVLAError("GR00T raw DiT conditioning input is not BF16 before explicit widen") - - base_embeddings = language_inputs["inputs_embeds"] - visual_pos_masks = language_inputs["visual_pos_masks"] - deepstack_visual_embeds = language_inputs["deepstack_visual_embeds"] - token_count = captures["qwen_inputs"]["input_ids"].shape[1] - if (base_embeddings.dtype != torch.bfloat16 or - tuple(base_embeddings.shape) != (1, token_count, EXPECTED_QWEN_HIDDEN_DIM) or - visual_pos_masks.dtype != torch.bool or - tuple(visual_pos_masks.shape) != (1, token_count) or - len(deepstack_visual_embeds) != 3): - raise StarVLAError("official Qwen prepared input layout changed") - visual_mask = visual_pos_masks[0] - visual_token_count = int(visual_mask.sum().item()) - prepared_embeddings = torch.zeros( - (token_count, 4, EXPECTED_QWEN_HIDDEN_DIM), - dtype=torch.bfloat16, - device=base_embeddings.device, - ) - prepared_embeddings[:, 0, :] = base_embeddings[0] - for index, value in enumerate(deepstack_visual_embeds): - if (value.dtype != torch.bfloat16 or - tuple(value.shape) != (visual_token_count, EXPECTED_QWEN_HIDDEN_DIM)): - raise StarVLAError("official Qwen DeepStack prepared input layout changed") - prepared_embeddings[visual_mask, index + 1, :] = value - if captures["action_inputs"][0].dtype != torch.bfloat16 or any( - value.dtype != torch.float32 for value in captures["action_inputs"][1:] - ): - raise StarVLAError("GR00T actions must become FP32 after the first Euler update") - - future = action_model.future_tokens.weight.unsqueeze(0) - position = action_model.position_embedding.weight[:EXPECTED_ACTION_HORIZON].unsqueeze(0) - for step in range(4): - dit_input = captures["dit_inputs"][step] - if tuple(dit_input.shape) != (1, EXPECTED_SEQUENCE_LENGTH, EXPECTED_DIT_WIDTH): - raise StarVLAError("official GR00T DiT query sequence shape changed") - if not torch.equal(dit_input[:, :EXPECTED_FUTURE_TOKEN_COUNT], future): - raise StarVLAError("official GR00T DiT query prefix is not future tokens") - expected_action_features = captures["action_encoder_outputs"][step] + position - if not torch.equal(dit_input[:, EXPECTED_FUTURE_TOKEN_COUNT:], expected_action_features): - raise StarVLAError("official GR00T DiT query suffix is not positioned action features") - if not torch.equal(captures["dit_conditioning_inputs"][step], captures["outer_raw_final"]): - raise StarVLAError("GR00T cross-attention conditioning changed across sampler steps") - if not torch.equal(captures["dit_outputs"][step], captures["decoder_inputs"][step]): - raise StarVLAError("GR00T DiT-to-velocity decoder boundary changed") - - velocities = [value[:, -EXPECTED_ACTION_HORIZON:] for value in captures["decoder_outputs"]] - actions_after = captures["action_inputs"][1:] + [captures["raw_policy"]] - previous = captures["initial_noise"].to(dtype=torch.float32) - for step, (velocity, actual) in enumerate(zip(velocities, actions_after)): - expected = previous + 0.25 * velocity - if not torch.equal(actual, expected): - raise StarVLAError(f"official GR00T Euler update mismatch at step {step}") - previous = actual - normalized = np.asarray(result.get("normalized_actions"), dtype=np.float32) - raw_policy, _ = _tensor_to_array(captures["raw_policy"]) - if normalized.shape != (1, 16, 7) or not np.array_equal(normalized, raw_policy): - raise StarVLAError("official normalized_actions differ from captured GR00T policy output") - if not np.isfinite(normalized).all(): - raise StarVLAError("official GR00T output contains non-finite values") - captures["block_outputs"] = block_outputs - captures["raw_qwen_taps"] = raw_qwen_taps - captures["prepared_embeddings"] = prepared_embeddings - captures["predicted_velocities"] = velocities - captures["actions_after_steps"] = actions_after - captures["normalized_actions"] = normalized - return captures - - -def _stack(values: Sequence[Any], *, label: str) -> tuple[np.ndarray, str]: - import torch - - if not values: - raise StarVLAError(f"cannot stack empty {label}") - dtype = values[0].dtype - if any(value.dtype != dtype for value in values): - raise StarVLAError(f"{label} has mixed source dtypes") - return _tensor_to_array(torch.stack(list(values), dim=0)) - - -def build_arrays(captures: Mapping[str, Any], unnormalized: np.ndarray) -> dict[str, np.ndarray]: - qwen_inputs = captures["qwen_inputs"] - input_ids = np.ascontiguousarray(qwen_inputs["input_ids"][0].cpu().numpy(), dtype=np.int64) - attention_mask = np.ascontiguousarray( - qwen_inputs["attention_mask"][0].to(dtype=__import__("torch").bool).cpu().numpy(), dtype=np.bool_ - ) - image_grid = np.ascontiguousarray(qwen_inputs["image_grid_thw"].cpu().numpy(), dtype=np.int64) - raw_final, _ = _tensor_to_array(captures["outer_raw_final"]) - initial_noise, _ = _tensor_to_array(captures["initial_noise"]) - dit_inputs, _ = _stack(captures["dit_inputs"], label="DiT inputs") - block_outputs, _ = _stack(captures["block_outputs"], label="DiT block outputs") - block_outputs = block_outputs.reshape( - 4, EXPECTED_DIT_BLOCK_COUNT, 1, EXPECTED_SEQUENCE_LENGTH, EXPECTED_DIT_WIDTH - ) - dit_outputs, _ = _stack(captures["dit_outputs"], label="DiT outputs") - velocities, _ = _stack(captures["predicted_velocities"], label="predicted velocities") - actions_after, _ = _stack(captures["actions_after_steps"], label="actions after steps") - action_features = np.ascontiguousarray( - dit_inputs[:, :, EXPECTED_FUTURE_TOKEN_COUNT:, :], dtype=np.float32 - ) - arrays = { - "input_ids": input_ids, - "attention_mask": attention_mask, - "image_grid_thw": image_grid, - "raw_l_out_35": np.ascontiguousarray(raw_final, dtype=np.float32), - "initial_noise": np.ascontiguousarray(initial_noise, dtype=np.float32), - "action_features": action_features, - "dit_inputs": np.ascontiguousarray(dit_inputs, dtype=np.float32), - "dit_block_outputs": np.ascontiguousarray(block_outputs, dtype=np.float32), - "dit_outputs": np.ascontiguousarray(dit_outputs, dtype=np.float32), - "predicted_velocities": np.ascontiguousarray(velocities, dtype=np.float32), - "actions_after_steps": np.ascontiguousarray(actions_after, dtype=np.float32), - "normalized_actions": np.ascontiguousarray(captures["normalized_actions"], dtype=np.float32), - "unnormalized_actions": np.ascontiguousarray(unnormalized, dtype=np.float32), - } - expected_shapes = { - "attention_mask": (input_ids.shape[0],), - "image_grid_thw": (1, 3), - "raw_l_out_35": (1, input_ids.shape[0], 2560), - "initial_noise": (1, 16, 7), - "action_features": (4, 1, 16, 768), - "dit_inputs": (4, 1, 48, 768), - "dit_block_outputs": (4, 16, 1, 48, 768), - "dit_outputs": (4, 1, 48, 1024), - "predicted_velocities": (4, 1, 16, 7), - "actions_after_steps": (4, 1, 16, 7), - "normalized_actions": (1, 16, 7), - "unnormalized_actions": (1, 16, 7), - } - for name, shape in expected_shapes.items(): - if arrays[name].shape != shape: - raise StarVLAError(f"official GR00T {name} shape mismatch: {arrays[name].shape}") - if any(not np.isfinite(value).all() for name, value in arrays.items() if name not in { - "input_ids", "attention_mask", "image_grid_thw" - }): - raise StarVLAError("official GR00T arrays contain non-finite values") - return arrays - - -def _runtime_record(torch: Any, transformers: Any, device: str, recorder: Mapping[str, Any]) -> dict[str, Any]: - cuda_device = torch.device(device) - index = cuda_device.index if cuda_device.index is not None else torch.cuda.current_device() - properties = torch.cuda.get_device_properties(index) - return { - "python": platform.python_version(), - "platform": platform.platform(), - "torch": torch.__version__, - "torchvision": _distribution_version("torchvision"), - "transformers": transformers.__version__, - "numpy": np.__version__, - "diffusers": _distribution_version("diffusers"), - "tokenizers": _distribution_version("tokenizers"), - "pillow": _distribution_version("Pillow"), - "omegaconf": _distribution_version("omegaconf"), - "accelerate": _distribution_version("accelerate"), - "safetensors": _distribution_version("safetensors"), - "official_environment_freeze": dict(OFFICIAL_ENVIRONMENT_FREEZE), - "cuda_runtime": torch.version.cuda, - "cudnn": torch.backends.cudnn.version(), - "device": str(cuda_device), - "device_name": properties.name, - "compute_capability": [properties.major, properties.minor], - "qwen3vl_recorder_probe": dict(recorder), - } - - -def _image_record(path: Path, image: Any) -> dict[str, Any]: - pixel_header = _canonical_json({"mode": image.mode, "size": list(image.size)}) - pixel_hash = hashlib.sha256(pixel_header + b"\x00" + image.tobytes()).hexdigest() - return { - "source_size": path.stat().st_size, - "source_sha256": sha256_file(path), - "decoded_mode": image.mode, - "decoded_size": list(image.size), - "decoded_pixel_sha256": pixel_hash, - } - - -def _load_images(paths: Iterable[Path]) -> tuple[list[Any], list[dict[str, Any]]]: - from PIL import Image - - images = [] - records = [] - for path in paths: - path = _regular_file(path, label="oracle image") - with Image.open(path) as opened: - opened.load() - image = opened.convert("RGB") - images.append(image) - records.append(_image_record(path, image)) - if len(images) != 1: - raise StarVLAError("official GR00T oracle requires exactly one image") - return images, records - - -def _write_runner_contract( - path: Path, *, golden_id: str, source: Mapping[str, Any], task: str, - model_instruction: str, unnorm_key: str, image_sha256: str, - array_records: Mapping[str, Any], token_count: int, initial_noise: np.ndarray, -) -> str: - raw_noise = np.ascontiguousarray(initial_noise, dtype=" Path: - import torch - import transformers - - output_dir = output_dir.resolve() - if output_dir.exists(): - raise StarVLAError(f"golden output directory already exists: {output_dir}") - output_dir.parent.mkdir(parents=True, exist_ok=True) - arrays = build_arrays(captures, unnormalized) - records = { - name: _array_record(value, source_dtype=ARRAY_SOURCE_DTYPES[name]) - for name, value in arrays.items() - } - model_instruction = expected_model_instruction(config, task) - identity = { - "schema_version": 1, - "variant": SUPPORTED_VARIANT, - "checkpoint_sha256": paths["variant"]["checkpoint"]["sha256"], - "starvla_revision": paths["catalog"]["source_revisions"]["starvla"], - "qwen_revision": paths["qwen"]["revision"], - "task": task, - "unnorm_key": unnorm_key, - "seed": SEED, - "images": [record["source_sha256"] for record in image_records], - } - golden_id = _sha256_bytes(_canonical_json(identity)) - variant = paths["variant"] - qwen = paths["qwen"] - source = { - "catalog_sha256": sha256_file(paths["catalog_path"]), - "bundle_uuid": official_bundle_uuid(variant, paths["catalog"]), - "starvla_repo_revision": paths["catalog"]["source_revisions"]["starvla"], - "checkpoint_repo_id": variant["repo_id"], - "checkpoint_revision": variant["revision"], - "checkpoint_size": variant["checkpoint"]["size"], - "checkpoint_sha256": variant["checkpoint"]["sha256"], - "policy_assets": _source_asset_hashes(variant), - "qwen_repo_id": qwen["repo_id"], - "qwen_revision": qwen["revision"], - "qwen_runtime_assets": _source_asset_hashes(qwen), - "qwen_converted_component_assets": _source_asset_hashes(qwen, staged=True), - "pinned_source_probe": paths["source_probe"], - } - with tempfile.TemporaryDirectory(prefix=f".{output_dir.name}.", dir=output_dir.parent) as temporary: - staging = Path(temporary) - inputs = staging / "inputs" - inputs.mkdir() - image_artifacts = [] - for index, (source_path, record) in enumerate(zip(image_paths, image_records)): - suffix = source_path.suffix.lower() or ".img" - destination = inputs / f"image-{index:02d}{suffix}" - shutil.copyfile(source_path, destination) - image_artifacts.append({ - **record, - "artifact": destination.relative_to(staging).as_posix(), - "artifact_size": destination.stat().st_size, - "artifact_sha256": sha256_file(destination), - }) - tensor_path = staging / "tensors.npz" - np.savez(tensor_path, **arrays) - contract_path = staging / "runner_contract.txt" - contract_sha = _write_runner_contract( - contract_path, golden_id=golden_id, source=source, task=task, - model_instruction=model_instruction, unnorm_key=unnorm_key, - image_sha256=image_records[0]["source_sha256"], array_records=records, - token_count=int(arrays["input_ids"].shape[0]), initial_noise=arrays["initial_noise"], - ) - noise_path = staging / "initial_noise.f32" - manifest: dict[str, Any] = { - "schema_version": GOLDEN_SCHEMA_VERSION, - "kind": GOLDEN_KIND, - "golden_id": golden_id, - "created_utc": dt.datetime.now(dt.timezone.utc).isoformat(), - "variant": SUPPORTED_VARIANT, - "model_type": "starvla", - "source": source, - "runtime": _runtime_record(torch, transformers, str(next(framework.parameters()).device), recorder_probe), - "determinism": { - "seed": SEED, - "rng_reset_immediately_before_predict": True, - "initial_noise_saved_explicitly": True, - "cross_language_seed_replay_allowed": False, - "torch_deterministic_algorithms": True, - "cublas_workspace_config": ":4096:8", - "cuda_matmul_allow_tf32": False, - "cudnn_allow_tf32": False, - "cudnn_benchmark": False, - "attention_implementation": "sdpa", - }, - "compatibility": { - "qwen_bootstrap": "config_only_then_strict_official_checkpoint_load", - "policy_parameters_after_strict_load": "float32", - "raw_qwen_conditioning": "bfloat16_complete_l_out_35", - "initial_noise": "torch_randn_bfloat16_then_exact_widen_at_action_encoder", - "dit_conditioning": "bfloat16_raw_l_out_35_exact_widen_to_float32", - "actions_dtype_by_step_input": ["bfloat16", "float32", "float32", "float32"], - "reason": ( - "pinned source requests CUDA autocast(dtype=float32), disabled by PyTorch 2.6; " - "the two explicit widens realize the declared FP32 action-policy path" - ), - }, - "input": { - "task": task, - "unnorm_key": unnorm_key, - "state": None, - "images": image_artifacts, - }, - "prompt": { - "framework_instruction": task, - "model_instruction": model_instruction, - "action_token_mode": "none", - }, - "model_contract": { - "action_horizon": 16, - "action_dim": 7, - "qwen_hidden_dim": 2560, - "qwen_tap": "outer_hidden_states_last_equals_complete_raw_l_out_35", - "attention_mask": "full_sequence_bool_nonzero_participates", - "future_token_count": 32, - "query_sequence_length": 48, - "query_token_order": "future_tokens_then_action_tokens", - "dit_width": 768, - "dit_output_dim": 1024, - "dit_block_count": 16, - "cross_attention_blocks": list(range(0, 16, 2)), - "self_attention_blocks": list(range(1, 16, 2)), - "timestep_ids": EXPECTED_TIMESTEP_IDS, - "euler_dt": 0.25, - "state_input_active": False, - "tap_layouts": { - "raw_l_out_35": "batch_token_hidden", - "action_features": "step_batch_action_width", - "dit_inputs": "step_batch_query_width", - "dit_block_outputs": "step_block_batch_query_width", - "dit_outputs": "step_batch_query_output", - "predicted_velocities": "step_batch_action_dimension", - "actions_after_steps": "step_batch_action_dimension", - }, - }, - "tokens": { - "input_ids": arrays["input_ids"].tolist(), - "attention_mask": arrays["attention_mask"].astype(np.uint8).tolist(), - "image_grid_thw": arrays["image_grid_thw"].tolist(), - }, - "outputs": { - "normalized_actions": arrays["normalized_actions"].tolist(), - "unnormalized_actions": arrays["unnormalized_actions"].tolist(), - }, - "artifacts": { - "tensors": { - "path": tensor_path.name, - "size": tensor_path.stat().st_size, - "sha256": sha256_file(tensor_path), - "encoding": "numpy_npz_stored", - "arrays": records, - }, - "initial_noise_raw": { - "path": noise_path.name, - "size": noise_path.stat().st_size, - "sha256": sha256_file(noise_path), - "encoding": "little_endian_float32_exact_widened_bfloat16", - "array_sha256": records["initial_noise"]["sha256"], - }, - "runner_contract": { - "path": contract_path.name, - "size": contract_path.stat().st_size, - "sha256": contract_sha, - "encoding": "ordered_utf8_key_value_v1", - }, - }, - } - manifest["integrity"] = { - "canonicalization": "utf8_json_sort_keys_compact_excluding_integrity", - "manifest_payload_sha256": _sha256_bytes(_canonical_json(manifest)), - } - (staging / "golden.json").write_text( - json.dumps(manifest, ensure_ascii=False, indent=2, sort_keys=True) + "\n", - encoding="utf-8", - ) - Path(temporary).replace(output_dir) - return output_dir / "golden.json" - - -def _preflight_record(paths: Mapping[str, Any], recorder: Mapping[str, Any]) -> dict[str, Any]: - config = resolve_effective_config(Path(paths["policy_dir"]), SUPPORTED_VARIANT) - _validate_effective_config(config) - return { - "variant": SUPPORTED_VARIANT, - "checkpoint": str(paths["checkpoint"]), - "checkpoint_ready": paths["checkpoint_ready"], - "expected_checkpoint_size": paths["variant"]["checkpoint"]["size"], - "expected_checkpoint_sha256": paths["variant"]["checkpoint"]["sha256"], - "source_probe": paths["source_probe"], - "qwen3vl_recorder_probe": recorder, - "effective_config_valid": True, - } - - -def build_parser() -> argparse.ArgumentParser: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--checkpoint-root", type=Path, default=Path("ckpts/starvla")) - parser.add_argument("--starvla-source", type=Path, default=Path("ckpts/starvla/source/starvla")) - parser.add_argument("--catalog", type=Path, default=DEFAULT_CATALOG) - parser.add_argument("--image", type=Path, action="append", default=[]) - parser.add_argument("--task", default="grab the block.") - parser.add_argument("--unnorm-key", default="oxe_bridge") - parser.add_argument("--device", default="cuda:0") - parser.add_argument("--output-dir", type=Path, default=Path("goldens/starvla/groot/bridge-grab-block")) - parser.add_argument( - "--qwen-layer-diagnostic", - type=Path, - help="optionally write 36 x token_count x 2560 raw decoder taps as little-endian FP32", - ) - parser.add_argument( - "--qwen-prepared-embeddings", - type=Path, - help=("optionally write token_count x 10240 prepared decoder inputs as " - "little-endian FP32 exact-widened BF16"), - ) - parser.add_argument("--preflight-only", action="store_true") - return parser - - -def main(argv: Sequence[str] | None = None) -> int: - args = build_parser().parse_args(argv) - try: - if not sys.flags.isolated: - raise StarVLAError( - "the GR00T oracle must run in isolated mode; invoke with `python -I`" - ) - import torch - import transformers - - validate_runtime_versions( - torch_version=torch.__version__, - torchvision_version=_distribution_version("torchvision"), - transformers_version=transformers.__version__, - numpy_version=np.__version__, - diffusers_version=_distribution_version("diffusers"), - tokenizers_version=_distribution_version("tokenizers"), - pillow_version=_distribution_version("Pillow"), - omegaconf_version=_distribution_version("omegaconf"), - accelerate_version=_distribution_version("accelerate"), - safetensors_version=_distribution_version("safetensors"), - ) - _configure_determinism(torch, seed=SEED, device=args.device) - paths = validate_available_inputs( - checkpoint_root=args.checkpoint_root, - source_dir=args.starvla_source, - catalog_path=args.catalog, - ) - recorder = verify_transformers_qwen3vl_recorder_semantics(torch, transformers) - if args.preflight_only: - print(json.dumps(_preflight_record(paths, recorder), indent=2, sort_keys=True)) - return 0 - if not paths["checkpoint_ready"]: - raise StarVLAError( - f"official GR00T checkpoint is not ready: {paths['checkpoint']}" - ) - if len(args.image) != 1: - raise StarVLAError("exactly one --image is required") - images, image_records = _load_images(args.image) - framework, config = load_official_framework(paths, device=args.device) - captures = run_official_forward(framework, images=images, task=args.task) - if args.qwen_layer_diagnostic is not None: - diagnostic_path = args.qwen_layer_diagnostic.resolve() - if not diagnostic_path.parent.is_dir() or diagnostic_path.exists(): - raise StarVLAError( - "Qwen layer diagnostic parent must exist and output must be absent: " - f"{diagnostic_path}" - ) - raw_qwen_taps, source_dtype = _stack( - captures["raw_qwen_taps"], label="raw Qwen layer taps" - ) - token_count = captures["qwen_inputs"]["input_ids"].shape[1] - expected_source_shape = (36, 1, token_count, 2560) - if (source_dtype != "bfloat16" or - raw_qwen_taps.shape != expected_source_shape): - raise StarVLAError( - "official raw Qwen layer diagnostic has an incompatible dtype or shape" - ) - raw_qwen_taps = raw_qwen_taps[:, 0] - with diagnostic_path.open("xb") as stream: - stream.write(np.ascontiguousarray(raw_qwen_taps, dtype=" bytes: - return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode("utf-8") - - -def _sha256_bytes(value: bytes) -> str: - return hashlib.sha256(value).hexdigest() - - -def _array_sha256(value: np.ndarray) -> str: - array = np.ascontiguousarray(value) - header = _canonical_json({"dtype": array.dtype.str, "shape": list(array.shape)}) - return _sha256_bytes(header + b"\x00" + array.tobytes(order="C")) - - -def _array_record(value: np.ndarray, *, source_dtype: str | None = None) -> dict[str, Any]: - array = np.ascontiguousarray(value) - record: dict[str, Any] = { - "dtype": array.dtype.str, - "shape": list(array.shape), - "sha256": _array_sha256(array), - } - if source_dtype is not None: - record["source_dtype"] = source_dtype - return record - - -def _distribution_version(name: str) -> str: - try: - return importlib.metadata.version(name) - except importlib.metadata.PackageNotFoundError: - return "missing" - - -def _base_version(version: str) -> str: - return version.split("+", 1)[0] - - -def validate_runtime_versions( - *, - torch_version: str, - torchvision_version: str, - transformers_version: str, - numpy_version: str, -) -> None: - expected = { - "torch": EXPECTED_TORCH_VERSION, - "torchvision": EXPECTED_TORCHVISION_VERSION, - "transformers": EXPECTED_TRANSFORMERS_VERSION, - "numpy": EXPECTED_NUMPY_VERSION, - } - actual = { - "torch": _base_version(torch_version), - "torchvision": _base_version(torchvision_version), - "transformers": _base_version(transformers_version), - "numpy": _base_version(numpy_version), - } - mismatches = [ - f"{name}: expected {expected[name]}, got {actual[name]}" - for name in expected - if actual[name] != expected[name] - ] - if mismatches: - raise StarVLAError("official oracle runtime version mismatch: " + "; ".join(mismatches)) - - -def select_action_positions( - input_ids: np.ndarray, - *, - action_token_id: int, - chunk_len: int, -) -> tuple[list[list[int]], list[list[int]]]: - ids = np.asarray(input_ids) - if ids.ndim != 2: - raise StarVLAError(f"input_ids must be rank 2, got {list(ids.shape)}") - if chunk_len <= 0: - raise StarVLAError(f"action chunk length must be positive, got {chunk_len}") - - all_positions: list[list[int]] = [] - selected_positions: list[list[int]] = [] - for batch_index, row in enumerate(ids): - positions = np.flatnonzero(row == action_token_id).astype(np.int64).tolist() - if len(positions) < chunk_len: - raise StarVLAError( - f"sample {batch_index} has {len(positions)} action tokens; expected at least {chunk_len}" - ) - all_positions.append(positions) - selected_positions.append(positions[-chunk_len:]) - return all_positions, selected_positions - - -def expected_framework_instruction(config: Mapping[str, Any], task: str, chunk_len: int) -> str: - if not isinstance(task, str) or not task.strip(): - raise StarVLAError("task must be a non-empty string") - try: - vla_data = config["datasets"]["vla_data"] - except (KeyError, TypeError) as exc: - raise StarVLAError("checkpoint config has no datasets.vla_data object") from exc - if not isinstance(vla_data, Mapping): - raise StarVLAError("checkpoint config datasets.vla_data must be an object") - action_tokens = ACTION_TOKEN * chunk_len - return task + f" Please predict the next {chunk_len} robot actions: {action_tokens}." - - -def expected_model_instruction(config: Mapping[str, Any], framework_instruction: str) -> str: - """Mirror QWen3.build_qwenvl_inputs after QwenOFT adds its action suffix.""" - - try: - vla_data = config["datasets"]["vla_data"] - except (KeyError, TypeError) as exc: - raise StarVLAError("checkpoint config has no datasets.vla_data object") from exc - if not isinstance(vla_data, Mapping): - raise StarVLAError("checkpoint config datasets.vla_data must be an object") - cot_prompt = vla_data.get("CoT_prompt") - return ( - cot_prompt.replace("{instruction}", framework_instruction) - if isinstance(cot_prompt, str) - else framework_instruction - ) - - -def _run_git(source_dir: Path, *arguments: str) -> str: - try: - result = subprocess.run( - ["git", "-C", str(source_dir), *arguments], - check=True, - capture_output=True, - text=True, - ) - except (OSError, subprocess.CalledProcessError) as exc: - raise StarVLAError(f"failed to inspect pinned StarVLA checkout {source_dir}: {exc}") from exc - return result.stdout.strip() - - -def verify_pinned_source_checkout(source_dir: Path, expected_revision: str) -> None: - source_dir = source_dir.resolve() - if not (source_dir / ".git").exists(): - raise StarVLAError(f"StarVLA source is not a Git checkout: {source_dir}") - actual_revision = _run_git(source_dir, "rev-parse", "HEAD") - if actual_revision != expected_revision: - raise StarVLAError( - f"StarVLA source revision mismatch: expected {expected_revision}, got {actual_revision}" - ) - changes = _run_git(source_dir, "status", "--porcelain=v1", "--untracked-files=all") - if changes: - raise StarVLAError(f"pinned StarVLA checkout has tracked or untracked changes:\n{changes}") - - -def _ensure_regular_file(path: Path, *, label: str) -> None: - if not path.is_file() or path.is_symlink(): - raise StarVLAError(f"{label} must be a regular, non-symlink file: {path}") - - -def validate_official_inputs( - *, - checkpoint_root: Path, - source_dir: Path, - catalog_path: Path = DEFAULT_CATALOG, -) -> dict[str, Any]: - catalog = load_catalog(catalog_path) - variant = get_variant(catalog, SUPPORTED_VARIANT) - qwen = catalog["shared_assets"]["qwen3_vl_4b_instruct"] - checkpoint_root = checkpoint_root.resolve() - policy_dir = checkpoint_root / "sources" / variant["directory"] / variant["revision"] - qwen_dir = checkpoint_root / "sources" / qwen["directory"] / qwen["revision"] - checkpoint = policy_dir / variant["checkpoint"]["path"] - - expected_source = (checkpoint_root / "source" / "starvla").resolve() - if source_dir.resolve() != expected_source: - raise StarVLAError( - f"StarVLA source must be the canonical checkout {expected_source}, got {source_dir.resolve()}" - ) - verify_pinned_source_checkout(source_dir, catalog["source_revisions"]["starvla"]) - verify_catalog_files(policy_dir, variant) - verify_catalog_files(qwen_dir, qwen) - _ensure_regular_file(checkpoint, label="official OFT checkpoint") - incomplete_sidecar = Path(f"{checkpoint}.aria2") - if incomplete_sidecar.exists(): - raise StarVLAError( - f"official OFT checkpoint download is incomplete ({incomplete_sidecar} exists); resume the download first" - ) - verify_checkpoint_file(checkpoint, variant) - return { - "catalog": catalog, - "variant": variant, - "qwen": qwen, - "policy_dir": policy_dir, - "qwen_dir": qwen_dir, - "checkpoint": checkpoint, - "source_dir": source_dir.resolve(), - "catalog_path": catalog_path.resolve(), - } - - -def _require_isolated_python() -> None: - if not sys.flags.isolated: - raise StarVLAError( - "the oracle must run in Python isolated mode; invoke it as `python -I " - "tools/hf2gguf/starvla/generate_starvla_oft_golden.py ...`" - ) - - -def _configure_determinism(torch: Any, *, seed: int, device: str) -> None: - os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" - if not device.startswith("cuda"): - raise StarVLAError("the official OFT golden oracle currently requires a CUDA device") - if not torch.cuda.is_available(): - raise StarVLAError("CUDA is not available to PyTorch") - try: - device_index = torch.device(device).index - except (RuntimeError, ValueError) as exc: - raise StarVLAError(f"invalid CUDA device {device!r}: {exc}") from exc - torch.cuda.set_device(0 if device_index is None else device_index) - if not torch.cuda.is_bf16_supported(): - raise StarVLAError(f"CUDA device {device!r} does not support bfloat16") - - np.random.seed(seed) - torch.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - torch.use_deterministic_algorithms(True) - torch.backends.cuda.matmul.allow_tf32 = False - torch.backends.cudnn.allow_tf32 = False - torch.backends.cudnn.benchmark = False - - -@contextlib.contextmanager -def _config_only_qwen_bootstrap(torch: Any, transformers: Any, qwen_dir: Path): - """Make the official wrapper construct Qwen topology without duplicate base weights. - - The subsequent strict load supplies every persistent parameter from the - SHA256-pinned StarVLA checkpoint. This is equivalent to StarVLA's released - loader after its bootstrap base weights are overwritten, while avoiding a - second 9 GB model download. - """ - - model_class = transformers.Qwen3VLForConditionalGeneration - had_local_override = "from_pretrained" in model_class.__dict__ - original_local_override = model_class.__dict__.get("from_pretrained") - - def from_config_only(model_id: str | os.PathLike[str], **kwargs: Any): - actual = Path(model_id).resolve() - if actual != qwen_dir.resolve(): - raise StarVLAError(f"official wrapper requested unexpected Qwen source: {actual}") - if kwargs.get("dtype") not in (None, torch.bfloat16): - raise StarVLAError(f"unexpected Qwen bootstrap dtype: {kwargs.get('dtype')!r}") - config = transformers.AutoConfig.from_pretrained( - actual, - local_files_only=True, - trust_remote_code=False, - ) - if getattr(config, "model_type", None) != "qwen3_vl": - raise StarVLAError(f"unexpected pinned Qwen model_type: {getattr(config, 'model_type', None)!r}") - previous_dtype = torch.get_default_dtype() - try: - torch.set_default_dtype(torch.bfloat16) - with transformers.modeling_utils.no_init_weights(): - model = model_class(config) - finally: - torch.set_default_dtype(previous_dtype) - return model - - model_class.from_pretrained = staticmethod(from_config_only) - try: - yield - finally: - if had_local_override: - model_class.from_pretrained = original_local_override - else: - delattr(model_class, "from_pretrained") - - -def _assert_module_origin(module: Any, source_dir: Path) -> None: - module_path = Path(module.__file__).resolve() - try: - module_path.relative_to(source_dir.resolve()) - except ValueError as exc: - raise StarVLAError(f"imported StarVLA module is outside the pinned checkout: {module_path}") from exc - - -@contextlib.contextmanager -def _official_qwen_model_alias(qwen_dir: Path): - """Expose the pinned local model under the case-sensitive name StarVLA dispatches on.""" - qwen_dir = qwen_dir.resolve() - with tempfile.TemporaryDirectory(prefix="starvla-qwen-alias-") as temporary: - alias = Path(temporary) / "Qwen3-VL-4B-Instruct" - alias.symlink_to(qwen_dir, target_is_directory=True) - if alias.resolve() != qwen_dir: - raise StarVLAError(f"temporary Qwen alias did not resolve to the pinned model: {alias}") - yield alias - - -def load_official_framework(paths: Mapping[str, Any], *, device: str) -> tuple[Any, dict[str, Any]]: - import torch - import transformers - - source_dir = Path(paths["source_dir"]) - if any(name == "starVLA" or name.startswith("starVLA.") for name in sys.modules): - raise StarVLAError("starVLA was imported before pinned-source verification") - sys.path.insert(0, str(source_dir)) - try: - from starVLA.model.framework import base_framework, share_tools - from starVLA.model.framework.VLM4A import QwenOFT - - _assert_module_origin(base_framework, source_dir) - _assert_module_origin(share_tools, source_dir) - _assert_module_origin(QwenOFT, source_dir) - config, norm_stats = share_tools.read_mode_config(str(paths["checkpoint"])) - qwen_dir = Path(paths["qwen_dir"]) - with _official_qwen_model_alias(qwen_dir) as qwen_alias: - config = base_framework.merge_config_overrides( - config, - [ - f"framework.qwenvl.base_vlm={qwen_alias}", - "framework.qwenvl.attn_implementation=sdpa", - ], - ) - expected_instruction = expected_framework_instruction( - config, - "contract probe", - int(config["framework"]["action_model"]["action_horizon"]), - ) - if ACTION_TOKEN * int(config["framework"]["action_model"]["action_horizon"]) not in expected_instruction: - raise StarVLAError("effective OFT config does not produce the pinned action-token contract") - - cfg = share_tools.dict_to_namespace(config) - cfg.trainer.pretrained_checkpoint = None - with _config_only_qwen_bootstrap(torch, transformers, qwen_dir): - framework = QwenOFT.Qwenvl_OFT(cfg) - - try: - state_dict = torch.load(paths["checkpoint"], map_location="cpu", weights_only=True) - except TypeError: - state_dict = torch.load(paths["checkpoint"], map_location="cpu") - if not isinstance(state_dict, Mapping) or not state_dict: - raise StarVLAError("official checkpoint did not contain a non-empty state_dict") - framework.load_state_dict(state_dict, strict=True) - del state_dict - gc.collect() - framework.norm_stats = norm_stats - - if type(framework).__name__ != "Qwenvl_OFT": - raise StarVLAError(f"unexpected official framework class: {type(framework).__name__}") - if int(framework.action_token_id) != ACTION_TOKEN_ID or framework.action_token != ACTION_TOKEN: - raise StarVLAError( - f"official action token mismatch: {framework.action_token!r}/{framework.action_token_id}" - ) - if int(framework.chunk_len) != 16: - raise StarVLAError(f"unexpected official OFT action horizon: {framework.chunk_len}") - qwen_dtypes = {parameter.dtype for parameter in framework.qwen_vl_interface.parameters()} - policy_dtypes = {parameter.dtype for parameter in framework.action_model.parameters()} - if qwen_dtypes != {torch.bfloat16}: - raise StarVLAError(f"unexpected Qwen parameter dtypes after strict load: {qwen_dtypes}") - if policy_dtypes != {torch.float32}: - raise StarVLAError(f"unexpected OFT parameter dtypes after strict load: {policy_dtypes}") - framework = framework.to(dtype=torch.bfloat16).to(device).eval() - if {parameter.dtype for parameter in framework.parameters()} != {torch.bfloat16}: - raise StarVLAError("official --use_bf16 cast did not cover the whole OFT model") - return framework, config - finally: - if sys.path and sys.path[0] == str(source_dir): - del sys.path[0] - - -def _tensor_to_array(tensor: Any) -> tuple[np.ndarray, str]: - source_dtype = str(tensor.dtype).removeprefix("torch.") - value = tensor.detach().cpu().contiguous() - if source_dtype == "bfloat16": - value = value.float() - return np.ascontiguousarray(value.numpy()), source_dtype - - -def _image_pixel_sha256(image: Any) -> str: - header = _canonical_json({"mode": image.mode, "size": list(image.size)}) - return _sha256_bytes(header + b"\x00" + image.tobytes()) - - -def _image_record(path: Path, image: Any) -> dict[str, Any]: - return { - "source_path": str(path.resolve()), - "source_size": path.stat().st_size, - "source_sha256": sha256_file(path), - "decoded_mode": image.mode, - "decoded_size": list(image.size), - "decoded_pixel_sha256": _image_pixel_sha256(image), - } - - -def _processed_image_records(images: Sequence[Any]) -> list[dict[str, Any]]: - return [ - { - "index": index, - "mode": image.mode, - "size": list(image.size), - "pixel_sha256": _image_pixel_sha256(image), - } - for index, image in enumerate(images) - ] - - -def run_official_forward(framework: Any, *, images: Sequence[Any], task: str) -> dict[str, Any]: - """Run Qwenvl_OFT.predict_action while capturing its real intermediate values.""" - - import torch - - captures: dict[str, Any] = {} - qwen = framework.qwen_vl_interface - action_model = framework.action_model - original_build = qwen.build_qwenvl_inputs - original_gather = framework._gather_action_token_embeddings - original_policy = action_model.predict_action - - def capture_build(*args: Any, **kwargs: Any): - batch_images = kwargs.get("images", args[0] if args else None) - instructions = kwargs.get("instructions", args[1] if len(args) > 1 else None) - captures["processed_images"] = list(batch_images[0]) - captures["framework_instructions"] = list(instructions) - result = original_build(*args, **kwargs) - captures["qwen_inputs"] = { - key: value.detach() - for key, value in result.items() - if isinstance(value, torch.Tensor) - } - return result - - def capture_gather(*args: Any, **kwargs: Any): - queries = original_gather(*args, **kwargs) - captures["action_queries_raw"] = queries.detach() - policy_dtype = next(action_model.parameters()).dtype - captures["policy_input_dtype"] = str(policy_dtype).removeprefix("torch.") - if queries.dtype != policy_dtype: - raise StarVLAError( - f"official whole-model BF16 dtype mismatch: queries={queries.dtype}, policy={policy_dtype}" - ) - return queries - - def capture_policy(*args: Any, **kwargs: Any): - captures["action_queries_policy"] = args[0].detach() - output = original_policy(*args, **kwargs) - captures["raw_policy"] = output.detach() - return output - - def capture_qwen_hidden(_module: Any, _inputs: Any, output: Any): - if not getattr(output, "hidden_states", None): - raise StarVLAError("official Qwen output did not include hidden_states") - captures["last_hidden_state"] = output.hidden_states[-1].detach() - - qwen.build_qwenvl_inputs = capture_build - framework._gather_action_token_embeddings = capture_gather - action_model.predict_action = capture_policy - hook = qwen.register_forward_hook(capture_qwen_hidden) - try: - result = framework.predict_action(examples=[{"image": list(images), "lang": task}]) - finally: - hook.remove() - qwen.build_qwenvl_inputs = original_build - framework._gather_action_token_embeddings = original_gather - action_model.predict_action = original_policy - - required = { - "processed_images", - "framework_instructions", - "qwen_inputs", - "action_queries_raw", - "action_queries_policy", - "raw_policy", - "last_hidden_state", - } - missing = sorted(required - set(captures)) - if missing: - raise StarVLAError(f"official OFT instrumentation did not capture: {missing}") - if "input_ids" not in captures["qwen_inputs"]: - raise StarVLAError("official Qwen preprocessing did not produce input_ids") - input_ids, _ = _tensor_to_array(captures["qwen_inputs"]["input_ids"]) - _, selected_positions = select_action_positions( - input_ids, - action_token_id=ACTION_TOKEN_ID, - chunk_len=int(framework.chunk_len), - ) - last_hidden = captures["last_hidden_state"] - positions = torch.as_tensor(selected_positions, device=last_hidden.device, dtype=torch.long) - expected_queries = last_hidden.gather( - 1, - positions.unsqueeze(-1).expand(-1, -1, last_hidden.shape[-1]), - ) - if not torch.equal(expected_queries, captures["action_queries_raw"]): - raise StarVLAError("captured action queries do not match final hidden state at selected token positions") - if captures["action_queries_raw"].dtype != torch.bfloat16: - raise StarVLAError(f"unexpected raw action-query dtype: {captures['action_queries_raw'].dtype}") - if captures["action_queries_policy"].dtype != torch.bfloat16: - raise StarVLAError(f"unexpected OFT input dtype: {captures['action_queries_policy'].dtype}") - if not torch.equal(captures["action_queries_raw"], captures["action_queries_policy"]): - raise StarVLAError("OFT policy input changed across the BF16 model boundary") - normalized = np.asarray(result.get("normalized_actions")) - raw_policy, _ = _tensor_to_array(captures["raw_policy"]) - expected_shape = (1, int(framework.chunk_len), int(action_model.action_dim)) - if normalized.shape != expected_shape: - raise StarVLAError(f"official OFT output shape mismatch: expected {expected_shape}, got {normalized.shape}") - if normalized.shape != raw_policy.shape or not np.array_equal(normalized, raw_policy): - raise StarVLAError("official normalized_actions differ from the raw OFT policy output") - if not np.isfinite(normalized).all(): - raise StarVLAError("official OFT policy produced NaN or infinite actions") - captures["normalized_actions"] = np.ascontiguousarray(normalized, dtype=np.float32) - return captures - - -def _render_model_prompt(framework: Any, processed_images: Sequence[Any], instruction: str) -> str: - messages = [ - { - "role": "user", - "content": [ - *({"type": "image", "image": image} for image in processed_images), - {"type": "text", "text": instruction}, - ], - } - ] - rendered = framework.qwen_vl_interface.processor.apply_chat_template( - messages, - tokenize=False, - add_generation_prompt=True, - ) - if not isinstance(rendered, str): - raise StarVLAError(f"official processor returned a non-string rendered prompt: {type(rendered)}") - return rendered - - -def _build_arrays(captures: Mapping[str, Any], unnormalized: np.ndarray) -> tuple[dict[str, np.ndarray], dict[str, Any]]: - arrays: dict[str, np.ndarray] = {} - records: dict[str, Any] = {} - - def add(name: str, value: Any) -> None: - if isinstance(value, np.ndarray): - array = np.ascontiguousarray(value) - source_dtype = None - else: - array, source_dtype = _tensor_to_array(value) - arrays[name] = array - records[name] = _array_record(array, source_dtype=source_dtype) - - for key, tensor in sorted(captures["qwen_inputs"].items()): - add(f"qwen_input__{key}", tensor) - add("last_hidden_state", captures["last_hidden_state"]) - add("action_queries_raw", captures["action_queries_raw"]) - add("action_queries_policy", captures["action_queries_policy"]) - add("raw_policy", captures["raw_policy"]) - add("normalized_actions", captures["normalized_actions"]) - add("unnormalized_actions", np.ascontiguousarray(unnormalized)) - return arrays, records - - -def _runtime_record(torch: Any, transformers: Any, device: str) -> dict[str, Any]: - cuda_device = torch.device(device) - index = cuda_device.index if cuda_device.index is not None else torch.cuda.current_device() - properties = torch.cuda.get_device_properties(index) - return { - "python": platform.python_version(), - "platform": platform.platform(), - "torch": torch.__version__, - "torchvision": _distribution_version("torchvision"), - "transformers": transformers.__version__, - "numpy": np.__version__, - "pillow": _distribution_version("Pillow"), - "omegaconf": _distribution_version("omegaconf"), - "cuda_runtime": torch.version.cuda, - "cudnn": torch.backends.cudnn.version(), - "device": str(cuda_device), - "device_name": properties.name, - "compute_capability": [properties.major, properties.minor], - } - - -def _copy_inputs(staging: Path, image_paths: Sequence[Path]) -> list[str]: - inputs_dir = staging / "inputs" - inputs_dir.mkdir() - relative_paths = [] - for index, source in enumerate(image_paths): - suffix = source.suffix.lower() if source.suffix else ".img" - destination = inputs_dir / f"image-{index:02d}{suffix}" - shutil.copyfile(source, destination) - relative_paths.append(destination.relative_to(staging).as_posix()) - return relative_paths - - -def write_golden( - *, - output_dir: Path, - paths: Mapping[str, Any], - framework: Any, - config: Mapping[str, Any], - image_paths: Sequence[Path], - source_image_records: Sequence[Mapping[str, Any]], - task: str, - unnorm_key: str, - captures: Mapping[str, Any], - unnormalized: np.ndarray, -) -> Path: - import torch - import transformers - - output_dir = output_dir.resolve() - if output_dir.exists(): - raise StarVLAError(f"golden output directory already exists: {output_dir}") - output_dir.parent.mkdir(parents=True, exist_ok=True) - arrays, array_records = _build_arrays(captures, unnormalized) - input_ids = arrays["qwen_input__input_ids"] - all_positions, selected_positions = select_action_positions( - input_ids, - action_token_id=ACTION_TOKEN_ID, - chunk_len=int(framework.chunk_len), - ) - expected_instruction = expected_framework_instruction(config, task, int(framework.chunk_len)) - captured_instructions = captures["framework_instructions"] - if captured_instructions != [expected_instruction]: - raise StarVLAError( - f"official prompt contract changed: expected {expected_instruction!r}, got {captured_instructions!r}" - ) - model_instruction = expected_model_instruction(config, expected_instruction) - rendered_prompt = _render_model_prompt(framework, captures["processed_images"], model_instruction) - token_strings = framework.qwen_vl_interface.processor.tokenizer.convert_ids_to_tokens(input_ids[0].tolist()) - - identity = { - "schema_version": GOLDEN_SCHEMA_VERSION, - "variant": SUPPORTED_VARIANT, - "checkpoint_sha256": paths["variant"]["checkpoint"]["sha256"], - "starvla_revision": paths["catalog"]["source_revisions"]["starvla"], - "qwen_revision": paths["qwen"]["revision"], - "task": task, - "unnorm_key": unnorm_key, - "images": [record["source_sha256"] for record in source_image_records], - } - golden_id = _sha256_bytes(_canonical_json(identity)) - - with tempfile.TemporaryDirectory(prefix=f".{output_dir.name}.", dir=output_dir.parent) as temporary: - staging = Path(temporary) - copied_images = _copy_inputs(staging, image_paths) - tensor_path = staging / "tensors.npz" - np.savez(tensor_path, **arrays) - - image_records = [] - for index, record in enumerate(source_image_records): - copied = staging / copied_images[index] - image_records.append( - { - **record, - "artifact": copied_images[index], - "artifact_size": copied.stat().st_size, - "artifact_sha256": sha256_file(copied), - } - ) - - raw_policy = arrays["raw_policy"].tolist() - normalized = arrays["normalized_actions"].tolist() - unnormalized_list = arrays["unnormalized_actions"].tolist() - manifest = { - "schema_version": GOLDEN_SCHEMA_VERSION, - "kind": "starvla_oft_official_python_oracle", - "golden_id": golden_id, - "created_utc": dt.datetime.now(dt.timezone.utc).isoformat(), - "variant": SUPPORTED_VARIANT, - "model_type": paths["variant"]["model_type"], - "source": { - "catalog": str(paths["catalog_path"]), - "catalog_sha256": sha256_file(paths["catalog_path"]), - "bundle_uuid": official_bundle_uuid(paths["variant"], paths["catalog"]), - "starvla_repo_revision": paths["catalog"]["source_revisions"]["starvla"], - "starvla_checkout": str(paths["source_dir"]), - "checkpoint_repo_id": paths["variant"]["repo_id"], - "checkpoint_revision": paths["variant"]["revision"], - "checkpoint_path": str(paths["checkpoint"]), - "checkpoint_size": paths["variant"]["checkpoint"]["size"], - "checkpoint_sha256": paths["variant"]["checkpoint"]["sha256"], - "qwen_repo_id": paths["qwen"]["repo_id"], - "qwen_revision": paths["qwen"]["revision"], - "config_json_sha256": paths["variant"]["file_hashes"]["config.json"]["sha256"], - "config_yaml_sha256": paths["variant"]["file_hashes"]["config.yaml"]["sha256"], - "dataset_statistics_sha256": paths["variant"]["file_hashes"]["dataset_statistics.json"]["sha256"], - }, - "runtime": _runtime_record(torch, transformers, str(next(framework.parameters()).device)), - "determinism": { - "seed": 0, - "torch_deterministic_algorithms": True, - "cublas_workspace_config": os.environ.get("CUBLAS_WORKSPACE_CONFIG"), - "allow_tf32": False, - "attention_implementation": "sdpa", - }, - "compatibility": { - "qwen_bootstrap": ( - "config-only topology construction; every persistent parameter is then populated by " - "strict loading of the pinned official checkpoint" - ), - "whole_model_cast": { - "to": "bfloat16", - "reason": "official Bridge server launch uses --use_bf16", - }, - }, - "input": { - "task": task, - "unnorm_key": unnorm_key, - "images": image_records, - "processed_images": _processed_image_records(captures["processed_images"]), - }, - "model_contract": { - "framework_class": f"{type(framework).__module__}.{type(framework).__name__}", - "action_token": ACTION_TOKEN, - "action_token_id": ACTION_TOKEN_ID, - "action_horizon": int(framework.chunk_len), - "action_dim": int(framework.action_model.action_dim), - "qwen_hidden_dim": int(framework.qwen_vl_interface.model.config.hidden_size), - }, - "prompt": { - "framework_instruction": expected_instruction, - "model_instruction": model_instruction, - "rendered_chat_template": rendered_prompt, - }, - "tokens": { - "input_ids": input_ids.tolist(), - "token_strings": token_strings, - "all_action_token_positions": all_positions, - "selected_action_token_positions": selected_positions, - }, - "outputs": { - "raw_policy": raw_policy, - "normalized_actions": normalized, - "unnormalized_actions": unnormalized_list, - }, - "artifacts": { - "tensors": { - "path": tensor_path.name, - "size": tensor_path.stat().st_size, - "sha256": sha256_file(tensor_path), - "arrays": array_records, - } - }, - } - manifest_path = staging / "golden.json" - manifest_path.write_text( - json.dumps(manifest, ensure_ascii=False, indent=2, sort_keys=True) + "\n", - encoding="utf-8", - ) - Path(temporary).replace(output_dir) - return output_dir / "golden.json" - - -def build_parser() -> argparse.ArgumentParser: - parser = argparse.ArgumentParser( - description="Generate an auditable golden from the pinned official StarVLA Qwen3-VL OFT checkpoint." - ) - parser.add_argument("--image", action="append", default=[], type=Path, help="Ordered image input; repeat for views") - parser.add_argument("--task", help="Robot task instruction") - parser.add_argument("--unnorm-key", choices=("oxe_bridge", "oxe_rt1")) - parser.add_argument("--output-dir", type=Path) - parser.add_argument("--checkpoint-root", type=Path, default=Path("ckpts/starvla")) - parser.add_argument("--starvla-source", type=Path, default=None) - parser.add_argument("--device", default="cuda:0") - parser.add_argument( - "--preflight-only", - action="store_true", - help="Verify pinned source/assets/checkpoint/runtime without allocating the model", - ) - return parser - - -def _load_images(image_paths: Iterable[Path]) -> tuple[list[Any], list[dict[str, Any]]]: - try: - from PIL import Image - except ImportError as exc: - raise StarVLAError("Pillow is required to load oracle images") from exc - images = [] - records = [] - for path in image_paths: - path = path.resolve() - _ensure_regular_file(path, label="input image") - try: - with Image.open(path) as opened: - opened.load() - image = opened.copy() - except (OSError, ValueError) as exc: - raise StarVLAError(f"failed to decode input image {path}: {exc}") from exc - images.append(image) - records.append(_image_record(path, image)) - return images, records - - -def main(argv: Sequence[str] | None = None) -> int: - args = build_parser().parse_args(argv) - _require_isolated_python() - if not args.preflight_only: - missing = [ - name - for name, value in ( - ("--image", args.image), - ("--task", args.task), - ("--unnorm-key", args.unnorm_key), - ("--output-dir", args.output_dir), - ) - if not value - ] - if missing: - raise StarVLAError("golden generation requires " + ", ".join(missing)) - if not args.task.strip(): - raise StarVLAError("--task must not be empty") - checkpoint_root = Path(args.checkpoint_root).resolve() - output_dir = args.output_dir.resolve() - if output_dir == checkpoint_root or checkpoint_root in output_dir.parents: - raise StarVLAError("--output-dir must not be inside the pinned checkpoint source tree") - - checkpoint_root = args.checkpoint_root.resolve() - source_dir = args.starvla_source or checkpoint_root / "source" / "starvla" - paths = validate_official_inputs( - checkpoint_root=checkpoint_root, - source_dir=source_dir, - catalog_path=DEFAULT_CATALOG, - ) - - try: - import torch - import transformers - except ImportError as exc: - raise StarVLAError(f"official StarVLA runtime dependency is missing: {exc}") from exc - validate_runtime_versions( - torch_version=torch.__version__, - torchvision_version=_distribution_version("torchvision"), - transformers_version=transformers.__version__, - numpy_version=np.__version__, - ) - _configure_determinism(torch, seed=0, device=args.device) - if args.preflight_only: - print("Pinned StarVLA OFT oracle preflight passed.") - return 0 - - images, source_image_records = _load_images(args.image) - framework, config = load_official_framework(paths, device=args.device) - captures = run_official_forward(framework, images=images, task=args.task) - - source_dir = Path(paths["source_dir"]) - sys.path.insert(0, str(source_dir)) - try: - from deployment.model_server import policy_norm_processor - - _assert_module_origin(policy_norm_processor, source_dir) - normalizer = policy_norm_processor.PolicyNormProcessor( - str(paths["checkpoint"]), - unnorm_key=args.unnorm_key, - ) - normalized = captures["normalized_actions"] - unnormalized = normalizer.unapply_actions(normalized[0])[None, ...] - if unnormalized.shape != normalized.shape or not np.isfinite(unnormalized).all(): - raise StarVLAError( - f"official action unnormalization returned invalid values/shape: {unnormalized.shape}" - ) - finally: - if sys.path and sys.path[0] == str(source_dir): - del sys.path[0] - - manifest = write_golden( - output_dir=args.output_dir, - paths=paths, - framework=framework, - config=config, - image_paths=args.image, - source_image_records=source_image_records, - task=args.task, - unnorm_key=args.unnorm_key, - captures=captures, - unnormalized=np.ascontiguousarray(unnormalized), - ) - print(f"Wrote official StarVLA OFT golden: {manifest}") - return 0 - - -if __name__ == "__main__": - try: - raise SystemExit(main()) - except StarVLAError as exc: - raise SystemExit(f"error: {exc}") from exc diff --git a/tools/hf2gguf/starvla/generate_starvla_pi_v3_golden.py b/tools/hf2gguf/starvla/generate_starvla_pi_v3_golden.py deleted file mode 100644 index 1bb0b91..0000000 --- a/tools/hf2gguf/starvla/generate_starvla_pi_v3_golden.py +++ /dev/null @@ -1,1831 +0,0 @@ -#!/usr/bin/env python3 -"""Generate an auditable oracle from the pinned official StarVLA PI_v3 checkpoint. - -The exporter executes the pinned StarVLA QwenPI_v3 implementation and records -both Transformers' effective outer-model conditioning tuple and cloned raw -decoder-layer outputs. This distinction matters in Transformers 4.57: -DeepStack updates the first three recorded decoder outputs in place, while the -outer conditional model retains the raw final decoder output. -""" - -from __future__ import annotations - -import argparse -import contextlib -import datetime as dt -import gc -import hashlib -import importlib.metadata -import json -import math -import os -import platform -import random -import shutil -import subprocess -import sys -import tempfile -from pathlib import Path -from types import MethodType -from typing import Any, Iterable, Mapping, Sequence - -import numpy as np - - -TOOLS_DIR = Path(__file__).resolve().parent -if str(TOOLS_DIR) not in sys.path: - sys.path.insert(0, str(TOOLS_DIR)) - -from starvla_checkpoint import ( # noqa: E402 - DEFAULT_CATALOG, - StarVLAError, - get_variant, - load_catalog, - official_bundle_uuid, - resolve_effective_config, - sha256_file, - verify_catalog_files, - verify_checkpoint_file, -) - - -GOLDEN_SCHEMA_VERSION = 1 -SUPPORTED_VARIANT = "pi_v3" -GOLDEN_KIND = "starvla_pi_v3_official_python_oracle" -SEED = 0 -EXPECTED_TRANSFORMERS_VERSION = "4.57.0" -EXPECTED_TORCH_VERSION = "2.6.0" -EXPECTED_TORCHVISION_VERSION = "0.21.0" -EXPECTED_NUMPY_VERSION = "1.26.4" -EXPECTED_DIFFUSERS_VERSION = "0.37.1" -EXPECTED_TOKENIZERS_VERSION = "0.22.2" -EXPECTED_PILLOW_VERSION = "12.1.1" -EXPECTED_OMEGACONF_VERSION = "2.3.0" -EXPECTED_ACCELERATE_VERSION = "1.5.2" -EXPECTED_SAFETENSORS_VERSION = "0.7.0" -EXPECTED_QWEN3VL_MODELING_SHA256 = "dd63ed3b124232735b3dca1bfa28f9d6b0d3f7182afcb75dde8f3e724b2b22da" -EXPECTED_TRANSFORMERS_GENERIC_SHA256 = "b117ffb2e9d513def41ce596eb82057b8e2811c6edf29ffd0bb634979240ebed" -EXPECTED_QWEN3VL_PROCESSING_SHA256 = "efd8d64aaf608aad1ffb3e6d503d6a99e5227d007df95c1d9fa905d998cda4a9" -EXPECTED_QWEN2VL_IMAGE_PROCESSING_FAST_SHA256 = ( - "09bfa9b17df7c3f0c6159bc34008ee50f21d2472cd5bae7e5c21ba1ca13a423c" -) -EXPECTED_QWEN2VL_IMAGE_PROCESSING_SHA256 = ( - "7820a0fcca107e75605e08d9b774285ca2b0316f857bc0225c779794705ecf4f" -) -OFFICIAL_ENVIRONMENT_FREEZE = { - "path": "wandb/wandb/run-20260426_011111-enstjn5q/files/requirements.txt", - "size": 4354, - "sha256": "de6b505238663ea8a218620e8a4f99cbcfe1e6e09f347ab26f68fe434f3fb00e", -} -EXPECTED_ACTION_HORIZON = 16 -EXPECTED_ACTION_DIM = 7 -EXPECTED_LAYER_COUNT = 36 -EXPECTED_QWEN_HIDDEN_DIM = 2560 -EXPECTED_PROJECTED_HIDDEN_DIM = 1024 -EXPECTED_TIMESTEP_IDS = [0, 250, 500, 750] -EXPECTED_COT_TEMPLATE = ( - "Your task is {instruction}. To identify the key objects for your task. " - "Locate their bounding boxes in [x1,y1,x2,y2] format." -) -CONDITIONING_TAP_NAMES = ( - [f"deepstack_out-{index}" for index in range(3)] - + [f"l_out-{index}" for index in range(3, EXPECTED_LAYER_COUNT)] -) -RAW_TAP_NAMES = [f"l_out-{index}" for index in range(EXPECTED_LAYER_COUNT)] -FINAL_NORM_DIAGNOSTIC_NAME = "result_norm" -CONDITIONING_SEMANTICS = ( - "Transformers 4.57 outer conditional-model recorder references after in-place DeepStack, " - "then raw decoder outputs including the final layer" -) -PROJECTOR_AUTOCAST_CONTRACT = { - "autocast_device_type": "cuda", - "autocast_dtype": "bfloat16", - "layer_norm_input_dtype": "bfloat16", - "layer_norm_parameter_dtype": "float32", - "layer_norm_compute_dtype": "float32", - "layer_norm_output_dtype": "float32", - "linear_input_operand_dtype": "bfloat16", - "linear_weight_operand_dtype": "bfloat16", - "linear_bias_operand_dtype": "bfloat16", - "linear_bias_application": "cublaslt_epilogue_bias", - "linear_operand_rounding": "round_to_nearest_even", - "linear_per_split_accumulator_dtype": "float32", - "linear_split_partial_storage_dtype": "bfloat16", - "linear_split_reduction_scheme": "output_type", - "allow_bf16_reduced_precision_reduction_setting_affects_gemm_and_bias": False, - "linear_output_dtype": "bfloat16", - "saved_output_dtype": "float32", - "saved_output_transport": "exact_widen_of_bfloat16_value", - "layer_norm_validation": "all_36_outputs_bitwise_equal_explicit_fp32_reconstruction", - "linear_validation": "all_36_outputs_bitwise_equal_explicit_bf16_operand_reconstruction", - "fp32_linear_then_output_round_is_distinct": True, -} - - -def _canonical_json(value: Any) -> bytes: - return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode("utf-8") - - -def _sha256_bytes(value: bytes) -> str: - return hashlib.sha256(value).hexdigest() - - -def _array_sha256(value: np.ndarray) -> str: - array = np.ascontiguousarray(value) - header = _canonical_json({"dtype": array.dtype.str, "shape": list(array.shape)}) - digest = hashlib.sha256() - digest.update(header) - digest.update(b"\x00") - payload = memoryview(array).cast("B") - for start in range(0, len(payload), 16 * 1024 * 1024): - digest.update(payload[start : start + 16 * 1024 * 1024]) - return digest.hexdigest() - - -def _array_record(value: np.ndarray, *, source_dtype: str | None = None) -> dict[str, Any]: - array = np.ascontiguousarray(value) - record: dict[str, Any] = { - "dtype": array.dtype.str, - "shape": list(array.shape), - "sha256": _array_sha256(array), - } - if source_dtype is not None: - record["source_dtype"] = source_dtype - return record - - -def _distribution_version(name: str) -> str: - try: - return importlib.metadata.version(name) - except importlib.metadata.PackageNotFoundError: - return "missing" - - -def _base_version(version: str) -> str: - return version.split("+", 1)[0] - - -def validate_runtime_versions( - *, - torch_version: str, - torchvision_version: str, - transformers_version: str, - numpy_version: str, - diffusers_version: str, - tokenizers_version: str, - pillow_version: str, - omegaconf_version: str, - accelerate_version: str, - safetensors_version: str, -) -> None: - expected = { - "torch": EXPECTED_TORCH_VERSION, - "torchvision": EXPECTED_TORCHVISION_VERSION, - "transformers": EXPECTED_TRANSFORMERS_VERSION, - "numpy": EXPECTED_NUMPY_VERSION, - "diffusers": EXPECTED_DIFFUSERS_VERSION, - "tokenizers": EXPECTED_TOKENIZERS_VERSION, - "pillow": EXPECTED_PILLOW_VERSION, - "omegaconf": EXPECTED_OMEGACONF_VERSION, - "accelerate": EXPECTED_ACCELERATE_VERSION, - "safetensors": EXPECTED_SAFETENSORS_VERSION, - } - actual = { - "torch": _base_version(torch_version), - "torchvision": _base_version(torchvision_version), - "transformers": _base_version(transformers_version), - "numpy": _base_version(numpy_version), - "diffusers": _base_version(diffusers_version), - "tokenizers": _base_version(tokenizers_version), - "pillow": _base_version(pillow_version), - "omegaconf": _base_version(omegaconf_version), - "accelerate": _base_version(accelerate_version), - "safetensors": _base_version(safetensors_version), - } - mismatches = [ - f"{name}: expected {expected[name]}, got {actual[name]}" - for name in expected - if actual[name] != expected[name] - ] - if mismatches: - raise StarVLAError("official PI_v3 oracle runtime version mismatch: " + "; ".join(mismatches)) - - -def expected_model_instruction(config: Mapping[str, Any], task: str) -> str: - if not isinstance(task, str) or not task.strip(): - raise StarVLAError("task must be a non-empty string") - try: - vla_data = config["datasets"]["vla_data"] - except (KeyError, TypeError) as exc: - raise StarVLAError("checkpoint config has no datasets.vla_data object") from exc - if not isinstance(vla_data, Mapping): - raise StarVLAError("checkpoint config datasets.vla_data must be an object") - cot_prompt = vla_data.get("CoT_prompt") - if not isinstance(cot_prompt, str) or cot_prompt.count("{instruction}") != 1: - raise StarVLAError("official PI_v3 CoT_prompt must contain exactly one {instruction} placeholder") - return cot_prompt.replace("{instruction}", task) - - -def expected_runtime_contract() -> dict[str, Any]: - """Describe the reference inputs needed to reproduce the action oracle.""" - return { - "conditioning": { - "hidden_tuple_indices": list(range(1, 37)), - "hidden_tap_names": CONDITIONING_TAP_NAMES, - }, - "timesteps": EXPECTED_TIMESTEP_IDS, - "action_shape": [16, 7], - } - - -def _run_git(source_dir: Path, *arguments: str) -> str: - try: - result = subprocess.run( - ["git", "-C", str(source_dir), *arguments], - check=True, - capture_output=True, - text=True, - ) - except (OSError, subprocess.CalledProcessError) as exc: - raise StarVLAError(f"failed to inspect pinned StarVLA checkout {source_dir}: {exc}") from exc - return result.stdout.strip() - - -def verify_pinned_source_checkout(source_dir: Path, expected_revision: str) -> None: - source_dir = source_dir.resolve() - if not (source_dir / ".git").exists(): - raise StarVLAError(f"StarVLA source is not a Git checkout: {source_dir}") - actual_revision = _run_git(source_dir, "rev-parse", "HEAD") - if actual_revision != expected_revision: - raise StarVLAError( - f"StarVLA source revision mismatch: expected {expected_revision}, got {actual_revision}" - ) - changes = _run_git(source_dir, "status", "--porcelain=v1", "--untracked-files=all") - if changes: - raise StarVLAError(f"pinned StarVLA checkout has tracked or untracked changes:\n{changes}") - - -def _ensure_regular_file(path: Path, *, label: str) -> None: - if not path.is_file() or path.is_symlink(): - raise StarVLAError(f"{label} must be a regular, non-symlink file: {path}") - - -def validate_official_inputs( - *, - checkpoint_root: Path, - source_dir: Path, - catalog_path: Path = DEFAULT_CATALOG, -) -> dict[str, Any]: - catalog = load_catalog(catalog_path) - variant = get_variant(catalog, SUPPORTED_VARIANT) - qwen = catalog["shared_assets"]["qwen3_vl_4b_instruct"] - checkpoint_root = checkpoint_root.resolve() - policy_dir = checkpoint_root / "sources" / variant["directory"] / variant["revision"] - qwen_dir = checkpoint_root / "sources" / qwen["directory"] / qwen["revision"] - checkpoint = policy_dir / variant["checkpoint"]["path"] - - expected_source = (checkpoint_root / "source" / "starvla").resolve() - if source_dir.resolve() != expected_source: - raise StarVLAError( - f"StarVLA source must be the canonical checkout {expected_source}, got {source_dir.resolve()}" - ) - verify_pinned_source_checkout(source_dir, catalog["source_revisions"]["starvla"]) - verify_catalog_files(policy_dir, variant) - verify_catalog_files(qwen_dir, qwen) - _ensure_regular_file(checkpoint, label="official PI_v3 checkpoint") - incomplete_sidecar = Path(f"{checkpoint}.aria2") - if incomplete_sidecar.exists(): - raise StarVLAError( - f"official PI_v3 checkpoint download is incomplete ({incomplete_sidecar} exists); resume it first" - ) - verify_checkpoint_file(checkpoint, variant) - return { - "catalog": catalog, - "variant": variant, - "qwen": qwen, - "policy_dir": policy_dir, - "qwen_dir": qwen_dir, - "checkpoint": checkpoint, - "source_dir": source_dir.resolve(), - "catalog_path": catalog_path.resolve(), - } - - -def _require_isolated_python() -> None: - if not sys.flags.isolated: - raise StarVLAError( - "the PI_v3 oracle must run in isolated mode; invoke it as `python -I " - "tools/hf2gguf/starvla/generate_starvla_pi_v3_golden.py ...`" - ) - - -def _configure_determinism(torch: Any, *, seed: int, device: str) -> None: - os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" - os.environ["NO_ALBUMENTATIONS_UPDATE"] = "1" - if not device.startswith("cuda"): - raise StarVLAError("the official PI_v3 golden oracle requires a CUDA device") - if not torch.cuda.is_available(): - raise StarVLAError("CUDA is not available to PyTorch") - try: - cuda_device = torch.device(device) - except (RuntimeError, ValueError) as exc: - raise StarVLAError(f"invalid CUDA device {device!r}: {exc}") from exc - torch.cuda.set_device(0 if cuda_device.index is None else cuda_device.index) - if not torch.cuda.is_bf16_supported(): - raise StarVLAError(f"CUDA device {device!r} does not support bfloat16") - - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - torch.use_deterministic_algorithms(True) - torch.backends.cuda.matmul.allow_tf32 = False - torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False - torch.backends.cudnn.allow_tf32 = False - torch.backends.cudnn.benchmark = False - - -def verify_transformers_qwen3vl_recorder_semantics(torch: Any, transformers: Any) -> dict[str, Any]: - """Execute the outer 4-layer Qwen3-VL model to gate 4.57 recorder semantics.""" - - try: - from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLConfig - from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLForConditionalGeneration - from transformers.models.qwen2_vl import ( - image_processing_qwen2_vl, - image_processing_qwen2_vl_fast, - ) - from transformers.models.qwen3_vl import modeling_qwen3_vl, processing_qwen3_vl - from transformers.utils import generic as transformers_generic - except ImportError as exc: - raise StarVLAError(f"Transformers lacks the pinned Qwen3-VL implementation: {exc}") from exc - modeling_path = Path(modeling_qwen3_vl.__file__).resolve() - actual_source_sha = sha256_file(modeling_path) - if actual_source_sha != EXPECTED_QWEN3VL_MODELING_SHA256: - raise StarVLAError( - "Transformers 4.57 Qwen3-VL implementation SHA256 mismatch: " - f"expected {EXPECTED_QWEN3VL_MODELING_SHA256}, got {actual_source_sha}" - ) - generic_path = Path(transformers_generic.__file__).resolve() - actual_generic_sha = sha256_file(generic_path) - if actual_generic_sha != EXPECTED_TRANSFORMERS_GENERIC_SHA256: - raise StarVLAError( - "Transformers 4.57 recorder implementation SHA256 mismatch: " - f"expected {EXPECTED_TRANSFORMERS_GENERIC_SHA256}, got {actual_generic_sha}" - ) - processing_path = Path(processing_qwen3_vl.__file__).resolve() - actual_processing_sha = sha256_file(processing_path) - if actual_processing_sha != EXPECTED_QWEN3VL_PROCESSING_SHA256: - raise StarVLAError( - "Transformers 4.57 Qwen3-VL processor implementation SHA256 mismatch: " - f"expected {EXPECTED_QWEN3VL_PROCESSING_SHA256}, got {actual_processing_sha}" - ) - image_processing_fast_path = Path(image_processing_qwen2_vl_fast.__file__).resolve() - actual_image_processing_fast_sha = sha256_file(image_processing_fast_path) - if actual_image_processing_fast_sha != EXPECTED_QWEN2VL_IMAGE_PROCESSING_FAST_SHA256: - raise StarVLAError( - "Transformers 4.57 Qwen2-VL fast image processor SHA256 mismatch: " - f"expected {EXPECTED_QWEN2VL_IMAGE_PROCESSING_FAST_SHA256}, " - f"got {actual_image_processing_fast_sha}" - ) - image_processing_path = Path(image_processing_qwen2_vl.__file__).resolve() - actual_image_processing_sha = sha256_file(image_processing_path) - if actual_image_processing_sha != EXPECTED_QWEN2VL_IMAGE_PROCESSING_SHA256: - raise StarVLAError( - "Transformers 4.57 Qwen2-VL smart-resize implementation SHA256 mismatch: " - f"expected {EXPECTED_QWEN2VL_IMAGE_PROCESSING_SHA256}, " - f"got {actual_image_processing_sha}" - ) - - config = Qwen3VLConfig( - text_config={ - "vocab_size": 32, - "hidden_size": 16, - "intermediate_size": 32, - "num_hidden_layers": 4, - "num_attention_heads": 2, - "num_key_value_heads": 1, - "head_dim": 8, - "max_position_embeddings": 32, - "use_cache": False, - "rope_scaling": { - "rope_type": "default", - "mrope_section": [2, 2, 4], - "mrope_interleaved": True, - }, - }, - vision_config={ - "depth": 1, - "hidden_size": 16, - "intermediate_size": 32, - "num_heads": 2, - "in_channels": 3, - "patch_size": 2, - "spatial_merge_size": 1, - "temporal_patch_size": 1, - "out_hidden_size": 16, - "num_position_embeddings": 16, - "deepstack_visual_indexes": [], - }, - image_token_id=2, - video_token_id=3, - vision_start_token_id=1, - vision_end_token_id=4, - ) - model = Qwen3VLForConditionalGeneration(config).cpu().eval() - raw: dict[int, Any] = {} - final_norm: dict[str, Any] = {} - inner_hidden: dict[str, Any] = {} - handles = [ - layer.register_forward_hook( - lambda _module, _inputs, output, index=index: raw.__setitem__( - index, output.detach().clone() - ) - ) - for index, layer in enumerate(model.model.language_model.layers) - ] - handles.append( - model.model.language_model.norm.register_forward_hook( - lambda _module, _inputs, output: final_norm.__setitem__( - "value", output.detach().clone() - ) - ) - ) - - def capture_inner(_module: Any, _inputs: Any, output: Any) -> None: - hidden_states = getattr(output, "hidden_states", None) - if hidden_states is not None: - inner_hidden["value"] = tuple(value.detach().clone() for value in hidden_states) - - handles.append(model.model.language_model.register_forward_hook(capture_inner)) - - def fake_get_image_features(_model: Any, pixel_values: Any, image_grid_thw: Any = None): - dtype = model.model.language_model.embed_tokens.weight.dtype - image_embed = torch.arange(16, dtype=dtype).reshape(1, 16) / 100.0 - deepstack = [ - torch.full((1, 16), float(index + 1), dtype=dtype) - for index in range(3) - ] - return [image_embed], deepstack - - model.model.get_image_features = MethodType(fake_get_image_features, model.model) - try: - input_ids = torch.tensor([[1, 2, 4, 5, 6]], dtype=torch.long) - visual_mask = input_ids == config.image_token_id - with torch.no_grad(): - output = model( - input_ids=input_ids, - attention_mask=torch.ones_like(input_ids), - pixel_values=torch.zeros(1), - image_grid_thw=torch.tensor([[1, 1, 1]], dtype=torch.long), - output_hidden_states=True, - use_cache=False, - logits_to_keep=1, - ) - finally: - for handle in handles: - handle.remove() - hidden = output.hidden_states - if hidden is None or len(hidden) != 5: - raise StarVLAError( - "Transformers 4.57 outer recorder probe did not return input + four layer states" - ) - for index in range(3): - delta = hidden[index + 1][visual_mask] - raw[index][visual_mask] - expected = torch.full_like(delta, float(index + 1)) - if not torch.allclose(delta, expected, rtol=0.0, atol=5e-7): - raise StarVLAError( - f"Transformers 4.57 recorder probe did not retain DeepStack's in-place layer {index} update" - ) - if not torch.equal(hidden[index + 1][~visual_mask], raw[index][~visual_mask]): - raise StarVLAError( - f"Transformers 4.57 recorder probe unexpectedly changed non-visual layer {index} tokens" - ) - if not torch.equal(hidden[-1], raw[3]): - raise StarVLAError( - "Transformers 4.57 outer recorder probe did not retain the raw final decoder output" - ) - if "value" not in final_norm: - raise StarVLAError("Transformers 4.57 recorder probe did not capture final RMSNorm") - if torch.equal(hidden[-1], final_norm["value"]): - raise StarVLAError( - "Transformers 4.57 outer recorder probe unexpectedly exposed result_norm as conditioning" - ) - inner = inner_hidden.get("value") - if inner is None or len(inner) != 5 or not torch.equal(inner[-1], final_norm["value"]): - raise StarVLAError( - "Transformers 4.57 inner recorder probe did not expose result_norm for diagnostics" - ) - return { - "modeling_qwen3_vl_path": str(modeling_path), - "modeling_qwen3_vl_sha256": actual_source_sha, - "transformers_generic_path": str(generic_path), - "transformers_generic_sha256": actual_generic_sha, - "processing_qwen3_vl_path": str(processing_path), - "processing_qwen3_vl_sha256": actual_processing_sha, - "image_processing_qwen2_vl_fast_path": str(image_processing_fast_path), - "image_processing_qwen2_vl_fast_sha256": actual_image_processing_fast_sha, - "image_processing_qwen2_vl_path": str(image_processing_path), - "image_processing_qwen2_vl_sha256": actual_image_processing_sha, - "model_class": "Qwen3VLForConditionalGeneration", - "observed_order": [ - "deepstack_out-0", - "deepstack_out-1", - "deepstack_out-2", - "l_out-3", - ], - "inner_terminal": "result_norm", - "outer_terminal": "l_out-3", - "mechanism": ( - "outer_recorder_keeps_raw_final_decoder_output_while_first_three_layer_references_receive_" - "in_place_deepstack_updates" - ), - } - - -@contextlib.contextmanager -def _config_only_qwen_bootstrap(torch: Any, transformers: Any, qwen_dir: Path): - model_class = transformers.Qwen3VLForConditionalGeneration - had_local_override = "from_pretrained" in model_class.__dict__ - original_local_override = model_class.__dict__.get("from_pretrained") - - def from_config_only(model_id: str | os.PathLike[str], **kwargs: Any): - actual = Path(model_id).resolve() - if actual != qwen_dir.resolve(): - raise StarVLAError(f"official wrapper requested unexpected Qwen source: {actual}") - if kwargs.get("dtype") not in (None, torch.bfloat16): - raise StarVLAError(f"unexpected Qwen bootstrap dtype: {kwargs.get('dtype')!r}") - config = transformers.AutoConfig.from_pretrained( - actual, - local_files_only=True, - trust_remote_code=False, - ) - if getattr(config, "model_type", None) != "qwen3_vl": - raise StarVLAError(f"unexpected pinned Qwen model_type: {getattr(config, 'model_type', None)!r}") - previous_dtype = torch.get_default_dtype() - try: - torch.set_default_dtype(torch.bfloat16) - with transformers.modeling_utils.no_init_weights(): - model = model_class(config) - finally: - torch.set_default_dtype(previous_dtype) - return model - - model_class.from_pretrained = staticmethod(from_config_only) - try: - yield - finally: - if had_local_override: - model_class.from_pretrained = original_local_override - else: - delattr(model_class, "from_pretrained") - - -@contextlib.contextmanager -def _official_qwen_model_alias(qwen_dir: Path): - """Expose the pinned local assets under StarVLA's case-sensitive dispatch name.""" - - qwen_dir = qwen_dir.resolve() - if not qwen_dir.is_dir(): - raise StarVLAError(f"pinned Qwen asset directory does not exist: {qwen_dir}") - with tempfile.TemporaryDirectory(prefix="starvla-qwen-alias-") as temporary: - alias = Path(temporary) / "Qwen3-VL-4B-Instruct" - alias.symlink_to(qwen_dir, target_is_directory=True) - if not alias.is_dir() or alias.resolve() != qwen_dir: - raise StarVLAError(f"temporary Qwen alias did not resolve to the pinned model: {alias}") - yield alias - - -def _assert_module_origin(module: Any, source_dir: Path) -> None: - module_path = Path(module.__file__).resolve() - try: - module_path.relative_to(source_dir.resolve()) - except ValueError as exc: - raise StarVLAError(f"imported StarVLA module is outside the pinned checkout: {module_path}") from exc - - -def verify_official_framework_import(paths: Mapping[str, Any]) -> dict[str, str]: - """Smoke-import the policy and normalizer from the already verified checkout.""" - source_dir = Path(paths["source_dir"]) - os.environ["NO_ALBUMENTATIONS_UPDATE"] = "1" - if any(name == "starVLA" or name.startswith("starVLA.") for name in sys.modules): - raise StarVLAError("starVLA was imported before pinned-source verification") - sys.path.insert(0, str(source_dir)) - try: - from deployment.model_server import policy_norm_processor - from starVLA.model.framework import base_framework, share_tools - from starVLA.model.framework.VLM4A import QwenPI_v3 - - modules = { - "base_framework": base_framework, - "share_tools": share_tools, - "qwen_pi_v3": QwenPI_v3, - "policy_norm_processor": policy_norm_processor, - } - for module in modules.values(): - _assert_module_origin(module, source_dir) - return {name: str(Path(module.__file__).resolve()) for name, module in modules.items()} - except ImportError as exc: - raise StarVLAError(f"failed to import the pinned official PI_v3 framework: {exc}") from exc - finally: - if sys.path and sys.path[0] == str(source_dir): - del sys.path[0] - - -def _validate_effective_config(config: Mapping[str, Any]) -> None: - try: - framework = config["framework"] - action = framework["action_model"] - diffusion = action["diffusion_model_cfg"] - vla_data = config["datasets"]["vla_data"] - except (KeyError, TypeError) as exc: - raise StarVLAError("effective PI_v3 config is missing required objects") from exc - actual = { - "framework": framework.get("name"), - "action_model_type": action.get("action_model_type"), - "action_horizon": action.get("action_horizon"), - "action_dim": action.get("action_dim"), - "state_dim": action.get("state_dim"), - "num_inference_timesteps": action.get("num_inference_timesteps"), - "num_timestep_buckets": action.get("num_timestep_buckets"), - "dit_width": diffusion.get("input_embedding_dim"), - "dit_layers": diffusion.get("num_layers"), - "interleave_self_attention": diffusion.get("interleave_self_attention"), - "use_canonical_forward": diffusion.get("use_canonical_forward"), - "image_size": vla_data.get("image_size"), - "data_mix": vla_data.get("data_mix"), - } - expected = { - "framework": "QwenPI_v3", - "action_model_type": "LayerwiseFM", - "action_horizon": 16, - "action_dim": 7, - "state_dim": 7, - "num_inference_timesteps": 4, - "num_timestep_buckets": 1000, - "dit_width": 1024, - "dit_layers": 36, - "interleave_self_attention": False, - "use_canonical_forward": True, - "image_size": [224, 224], - "data_mix": "bridge_rt_1", - } - if actual != expected: - raise StarVLAError(f"unexpected effective official PI_v3 config: {actual}") - if vla_data.get("CoT_prompt") != EXPECTED_COT_TEMPLATE: - raise StarVLAError(f"unexpected official PI_v3 CoT prompt: {vla_data.get('CoT_prompt')!r}") - expected_model_instruction(config, "contract probe") - - -def load_official_framework(paths: Mapping[str, Any], *, device: str) -> tuple[Any, dict[str, Any]]: - import torch - import transformers - - source_dir = Path(paths["source_dir"]) - if any(name == "starVLA" or name.startswith("starVLA.") for name in sys.modules): - raise StarVLAError("starVLA was imported before pinned-source verification") - sys.path.insert(0, str(source_dir)) - try: - from starVLA.model.framework import base_framework, share_tools - from starVLA.model.framework.VLM4A import QwenPI_v3 - - _assert_module_origin(base_framework, source_dir) - _assert_module_origin(share_tools, source_dir) - _assert_module_origin(QwenPI_v3, source_dir) - config = resolve_effective_config(Path(paths["policy_dir"]), SUPPORTED_VARIANT) - _validate_effective_config(config) - qwen_dir = Path(paths["qwen_dir"]).resolve() - with _official_qwen_model_alias(qwen_dir) as qwen_alias: - config = base_framework.merge_config_overrides( - config, - [ - f"framework.qwenvl.base_vlm={qwen_alias}", - "framework.qwenvl.attn_implementation=sdpa", - ], - ) - configured_qwen = Path(config["framework"]["qwenvl"]["base_vlm"]) - if "Qwen3-VL" not in str(configured_qwen) or configured_qwen.resolve() != qwen_dir: - raise StarVLAError( - "effective PI_v3 Qwen source does not preserve official dispatch and pinned assets" - ) - cfg = share_tools.dict_to_namespace(config) - cfg.trainer.pretrained_checkpoint = None - with _config_only_qwen_bootstrap(torch, transformers, qwen_dir): - framework = QwenPI_v3.Qwen_PI_v3(cfg) - - try: - state_dict = torch.load( - paths["checkpoint"], map_location="cpu", mmap=True, weights_only=True - ) - except TypeError: - state_dict = torch.load(paths["checkpoint"], map_location="cpu", weights_only=True) - if not isinstance(state_dict, Mapping) or not state_dict: - raise StarVLAError("official checkpoint did not contain a non-empty state_dict") - framework.load_state_dict(state_dict, strict=True) - del state_dict - gc.collect() - - if type(framework).__name__ != "Qwen_PI_v3": - raise StarVLAError(f"unexpected official framework class: {type(framework).__name__}") - action_model = framework.action_model - if int(framework.action_horizon) != EXPECTED_ACTION_HORIZON: - raise StarVLAError(f"unexpected official PI_v3 action horizon: {framework.action_horizon}") - if int(action_model.action_dim) != EXPECTED_ACTION_DIM: - raise StarVLAError(f"unexpected official PI_v3 action dimension: {action_model.action_dim}") - if len(framework.project_layers) != EXPECTED_LAYER_COUNT: - raise StarVLAError("official PI_v3 projector count is not 36") - if len(action_model.model.transformer_blocks) != EXPECTED_LAYER_COUNT: - raise StarVLAError("official PI_v3 DiT block count is not 36") - - qwen_dtypes = {parameter.dtype for parameter in framework.qwen_vl_interface.parameters()} - policy_dtypes = {parameter.dtype for parameter in action_model.parameters()} - projector_dtypes = {parameter.dtype for parameter in framework.project_layers.parameters()} - if qwen_dtypes != {torch.bfloat16}: - raise StarVLAError(f"unexpected Qwen parameter dtypes after strict load: {qwen_dtypes}") - if policy_dtypes != {torch.float32} or projector_dtypes != {torch.float32}: - raise StarVLAError( - "official PI_v3 FP32 policy/projector compatibility baseline changed: " - f"policy={policy_dtypes}, projectors={projector_dtypes}" - ) - return framework.to(device).eval(), config - finally: - if sys.path and sys.path[0] == str(source_dir): - del sys.path[0] - - -def _tensor_to_array(tensor: Any) -> tuple[np.ndarray, str]: - source_dtype = str(tensor.dtype).removeprefix("torch.") - value = tensor.detach().cpu().contiguous() - if source_dtype == "bfloat16": - value = value.float() - return np.ascontiguousarray(value.numpy()), source_dtype - - -def _require_tensor_equal(torch: Any, actual: Any, expected: Any, label: str) -> None: - if actual.shape != expected.shape or actual.dtype != expected.dtype or not torch.equal(actual, expected): - raise StarVLAError(f"official PI_v3 instrumentation mismatch for {label}") - - -def _projector_linear_bf16_operands( - torch: Any, - value: Any, - weight: Any, - bias: Any, -) -> Any: - """Replay CUDA autocast's Linear policy with explicit BF16 operands.""" - - with torch.autocast(value.device.type, enabled=False): - output = torch.nn.functional.linear( - value.to(dtype=torch.bfloat16), - weight.to(dtype=torch.bfloat16), - None if bias is None else bias.to(dtype=torch.bfloat16), - ) - if output.dtype != torch.bfloat16: - raise StarVLAError("explicit PI_v3 projector BF16 replay did not produce BF16") - return output - - -def _projector_linear_fp32_then_bf16( - torch: Any, - value: Any, - weight: Any, - bias: Any, -) -> Any: - """Represent the rejected FP32-Linear-then-BF16-round interpretation.""" - - with torch.autocast(value.device.type, enabled=False): - return torch.nn.functional.linear( - value.to(dtype=torch.float32), - weight.to(dtype=torch.float32), - None if bias is None else bias.to(dtype=torch.float32), - ).to(dtype=torch.bfloat16) - - -def run_official_forward( - framework: Any, - *, - images: Sequence[Any], - task: str, - seed: int = SEED, -) -> dict[str, Any]: - """Execute Qwen_PI_v3.predict_action and capture every parity boundary.""" - - import torch - - if torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction: - raise StarVLAError( - "PI_v3 official forward requires BF16 GEMM reduced-precision reduction to be disabled" - ) - captures: dict[str, Any] = {} - qwen = framework.qwen_vl_interface - action_model = framework.action_model - language_model = qwen.model.model.language_model - original_build = qwen.build_qwenvl_inputs - original_project = framework._project_vl_hidden_for_action - original_policy = action_model.predict_action - original_action_encoder = action_model.action_encoder.forward - original_dit = action_model.model.forward - raw_layer_outputs: dict[int, Any] = {} - deepstack_outputs: dict[int, Any] = {} - final_norm: dict[str, Any] = {} - projector_norm_inputs: dict[int, Any] = {} - projector_norm_outputs: dict[int, Any] = {} - projector_linear_inputs: dict[int, Any] = {} - projector_linear_outputs: dict[int, Any] = {} - projector_autocast_states: dict[tuple[int, str], tuple[bool, Any]] = {} - handles = [] - - def record_projector_tensor( - storage: dict[int, Any], - index: int, - value: Any, - label: str, - ) -> None: - if index in storage or not isinstance(value, torch.Tensor): - raise StarVLAError(f"official PI_v3 projector hook mismatch for {label} {index}") - storage[index] = value.detach() - - def record_projector_autocast(index: int, stage: str) -> None: - key = (index, stage) - if key in projector_autocast_states: - raise StarVLAError( - f"official PI_v3 projector autocast hook ran twice for {stage} {index}" - ) - projector_autocast_states[key] = ( - torch.is_autocast_enabled("cuda"), - torch.get_autocast_dtype("cuda"), - ) - - def capture_build(*args: Any, **kwargs: Any): - if "qwen_inputs" in captures: - raise StarVLAError("official PI_v3 preprocessing ran more than once") - batch_images = kwargs.get("images", args[0] if args else None) - instructions = kwargs.get("instructions", args[1] if len(args) > 1 else None) - captures["processed_images"] = list(batch_images[0]) - captures["framework_instructions"] = list(instructions) - result = original_build(*args, **kwargs) - captures["qwen_inputs"] = { - key: value.detach() for key, value in result.items() if isinstance(value, torch.Tensor) - } - return result - - def capture_project(hidden_states: Sequence[Any]): - if "project_input_taps" in captures: - raise StarVLAError("official PI_v3 projector bridge ran more than once") - captures["project_input_taps"] = [value.detach() for value in hidden_states] - projector_handles = [] - for index, projector in enumerate(framework.project_layers): - if ( - not isinstance(projector, torch.nn.Sequential) - or len(projector) != 2 - or not isinstance(projector[0], torch.nn.LayerNorm) - or not isinstance(projector[1], torch.nn.Linear) - ): - raise StarVLAError( - f"official PI_v3 projector {index} is no longer LayerNorm then Linear" - ) - - def capture_norm_input(_module: Any, inputs: Any, *, index: int = index) -> None: - if len(inputs) != 1: - raise StarVLAError(f"official PI_v3 projector norm {index} input arity changed") - record_projector_autocast(index, "layer_norm") - record_projector_tensor( - projector_norm_inputs, index, inputs[0], "LayerNorm input" - ) - - def capture_norm_output( - _module: Any, _inputs: Any, output: Any, *, index: int = index - ) -> None: - record_projector_tensor( - projector_norm_outputs, index, output, "LayerNorm output" - ) - - def capture_linear_input(_module: Any, inputs: Any, *, index: int = index) -> None: - if len(inputs) != 1: - raise StarVLAError(f"official PI_v3 projector linear {index} input arity changed") - record_projector_autocast(index, "linear") - record_projector_tensor( - projector_linear_inputs, index, inputs[0], "Linear logical input" - ) - - def capture_linear_output( - _module: Any, _inputs: Any, output: Any, *, index: int = index - ) -> None: - record_projector_tensor( - projector_linear_outputs, index, output, "Linear output" - ) - - projector_handles.extend( - [ - projector[0].register_forward_pre_hook(capture_norm_input), - projector[0].register_forward_hook(capture_norm_output), - projector[1].register_forward_pre_hook(capture_linear_input), - projector[1].register_forward_hook(capture_linear_output), - ] - ) - try: - projected = original_project(hidden_states) - finally: - for handle in projector_handles: - handle.remove() - captures["projected_hidden_taps"] = [value.detach() for value in projected] - return projected - - def capture_action_encoder(actions: Any, timesteps: Any): - return original_action_encoder(actions.to(dtype=torch.float32), timesteps) - - def capture_dit(*args: Any, **kwargs: Any): - conditioning = kwargs.get("encoder_hidden_states") - if not isinstance(conditioning, (list, tuple)): - raise StarVLAError("official PI_v3 DiT did not receive layer-wise conditioning") - kwargs["encoder_hidden_states"] = [value.to(dtype=torch.float32) for value in conditioning] - timestep = kwargs.get("timestep") - if timestep is None or timestep.numel() != 1: - raise StarVLAError("official PI_v3 DiT timestep shape changed") - captures.setdefault("timestep_ids", []).append(int(timestep.item())) - return original_dit(*args, **kwargs) - - def capture_policy(*args: Any, **kwargs: Any): - if "policy_input_taps" in captures: - raise StarVLAError("official PI_v3 policy sampler ran more than once") - policy_hidden = args[0] if args else kwargs.get("vl_embs_list") - if not isinstance(policy_hidden, (list, tuple)): - raise StarVLAError("official PI_v3 policy did not receive layer-wise hidden states") - captures["policy_input_taps"] = [value.detach() for value in policy_hidden] - original_randn = torch.randn - - def capture_randn(*randn_args: Any, **randn_kwargs: Any): - value = original_randn(*randn_args, **randn_kwargs) - if "initial_noise" in captures: - raise StarVLAError("official PI_v3 policy sampled noise more than once") - captures["initial_noise"] = value.detach().clone() - return value - - torch.randn = capture_randn - try: - output = original_policy(*args, **kwargs) - finally: - torch.randn = original_randn - captures["raw_policy"] = output.detach() - return output - - def capture_qwen_hidden(_module: Any, _inputs: Any, output: Any): - if "qwen_hidden_tuple" in captures: - raise StarVLAError("official PI_v3 outer Qwen model ran more than once") - hidden_states = getattr(output, "hidden_states", None) - if hidden_states is None or len(hidden_states) != EXPECTED_LAYER_COUNT + 1: - raise StarVLAError( - "official Transformers Qwen output did not contain input + 36 hidden states" - ) - captures["qwen_hidden_tuple"] = [value.detach() for value in hidden_states] - captures["conditioning_taps"] = [value.detach() for value in hidden_states[-36:]] - - for index, layer in enumerate(language_model.layers): - handles.append( - layer.register_forward_hook( - lambda _module, _inputs, output, index=index: raw_layer_outputs.__setitem__( - index, output.detach().clone() - ) - ) - ) - if index in (1, 2, 3): - handles.append( - layer.register_forward_pre_hook( - lambda _module, inputs, index=index: deepstack_outputs.__setitem__( - index - 1, inputs[0].detach().clone() - ) - ) - ) - handles.append( - language_model.norm.register_forward_hook( - lambda _module, _inputs, output: final_norm.__setitem__("value", output.detach().clone()) - ) - ) - handles.append(qwen.model.register_forward_hook(capture_qwen_hidden)) - qwen.build_qwenvl_inputs = capture_build - framework._project_vl_hidden_for_action = capture_project - action_model.predict_action = capture_policy - action_model.action_encoder.forward = capture_action_encoder - action_model.model.forward = capture_dit - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - try: - result = framework.predict_action(examples=[{"image": list(images), "lang": task}]) - finally: - for handle in handles: - handle.remove() - qwen.build_qwenvl_inputs = original_build - framework._project_vl_hidden_for_action = original_project - action_model.predict_action = original_policy - action_model.action_encoder.forward = original_action_encoder - action_model.model.forward = original_dit - - required = { - "processed_images", - "framework_instructions", - "qwen_inputs", - "qwen_hidden_tuple", - "conditioning_taps", - "project_input_taps", - "projected_hidden_taps", - "policy_input_taps", - "initial_noise", - "raw_policy", - "timestep_ids", - } - missing = sorted(required - set(captures)) - if missing: - raise StarVLAError(f"official PI_v3 instrumentation did not capture: {missing}") - if set(raw_layer_outputs) != set(range(EXPECTED_LAYER_COUNT)): - raise StarVLAError("official PI_v3 instrumentation missed raw Qwen decoder outputs") - if set(deepstack_outputs) != {0, 1, 2} or "value" not in final_norm: - raise StarVLAError("official PI_v3 instrumentation missed DeepStack/result_norm outputs") - for name in ( - "conditioning_taps", - "project_input_taps", - "projected_hidden_taps", - "policy_input_taps", - ): - if len(captures[name]) != EXPECTED_LAYER_COUNT: - raise StarVLAError(f"official PI_v3 {name} count is not 36") - - conditioning = captures["conditioning_taps"] - for index in range(3): - _require_tensor_equal(torch, conditioning[index], deepstack_outputs[index], CONDITIONING_TAP_NAMES[index]) - if torch.equal(conditioning[index], raw_layer_outputs[index]): - raise StarVLAError(f"DeepStack layer {index} did not change any recorded hidden-state value") - for index in range(3, EXPECTED_LAYER_COUNT): - _require_tensor_equal(torch, conditioning[index], raw_layer_outputs[index], f"l_out-{index}") - if torch.equal(conditioning[-1], final_norm["value"]): - raise StarVLAError("official outer Qwen conditioning unexpectedly ends at result_norm") - for actual, expected, name in zip(captures["project_input_taps"], conditioning, CONDITIONING_TAP_NAMES): - _require_tensor_equal(torch, actual, expected, f"projector input {name}") - for actual, expected in zip(captures["policy_input_taps"], captures["projected_hidden_taps"]): - _require_tensor_equal(torch, actual, expected, "projector-to-policy BF16 boundary") - - expected_projectors = set(range(EXPECTED_LAYER_COUNT)) - if any( - set(storage) != expected_projectors - for storage in ( - projector_norm_inputs, - projector_norm_outputs, - projector_linear_inputs, - projector_linear_outputs, - ) - ): - raise StarVLAError("official PI_v3 instrumentation missed a projector numeric boundary") - expected_autocast_keys = { - (index, stage) - for index in range(EXPECTED_LAYER_COUNT) - for stage in ("layer_norm", "linear") - } - if set(projector_autocast_states) != expected_autocast_keys: - raise StarVLAError("official PI_v3 instrumentation missed a projector autocast state") - - fp32_interpretation_is_distinct = False - with torch.inference_mode(): - for index, projector in enumerate(framework.project_layers): - norm_input = projector_norm_inputs[index] - norm_output = projector_norm_outputs[index] - linear_input = projector_linear_inputs[index] - linear_output = projector_linear_outputs[index] - projected_output = captures["projected_hidden_taps"][index] - if projector_autocast_states[(index, "layer_norm")] != (True, torch.bfloat16): - raise StarVLAError( - f"official PI_v3 projector {index} LayerNorm did not run under CUDA BF16 autocast" - ) - if projector_autocast_states[(index, "linear")] != (True, torch.bfloat16): - raise StarVLAError( - f"official PI_v3 projector {index} Linear did not run under CUDA BF16 autocast" - ) - if norm_input.dtype != torch.bfloat16: - raise StarVLAError(f"official PI_v3 projector {index} LayerNorm input is not BF16") - if norm_output.dtype != torch.float32 or linear_input.dtype != torch.float32: - raise StarVLAError( - f"official PI_v3 projector {index} LayerNorm did not expose an FP32 output" - ) - if linear_output.dtype != torch.bfloat16 or projected_output.dtype != torch.bfloat16: - raise StarVLAError(f"official PI_v3 projector {index} Linear output is not BF16") - _require_tensor_equal( - torch, norm_input, captures["project_input_taps"][index], f"projector {index} norm input" - ) - _require_tensor_equal( - torch, linear_input, norm_output, f"projector {index} norm-to-linear input" - ) - _require_tensor_equal( - torch, linear_output, projected_output, f"projector {index} linear output" - ) - - norm = projector[0] - linear = projector[1] - if ( - norm.weight is None - or norm.bias is None - or norm.weight.dtype != torch.float32 - or norm.bias.dtype != torch.float32 - or linear.weight.dtype != torch.float32 - or linear.bias is None - or linear.bias.dtype != torch.float32 - ): - raise StarVLAError( - f"official PI_v3 projector {index} FP32 parameter boundary changed" - ) - with torch.autocast(norm_input.device.type, enabled=False): - explicit_norm = torch.nn.functional.layer_norm( - norm_input.to(dtype=torch.float32), - norm.normalized_shape, - norm.weight, - norm.bias, - norm.eps, - ) - _require_tensor_equal( - torch, - explicit_norm, - norm_output, - f"projector {index} explicit FP32 LayerNorm reconstruction", - ) - explicit_bf16 = _projector_linear_bf16_operands( - torch, linear_input, linear.weight.detach(), linear.bias.detach() - ) - _require_tensor_equal( - torch, - explicit_bf16, - projected_output, - f"projector {index} explicit BF16 operand reconstruction", - ) - if not fp32_interpretation_is_distinct: - fp32_then_bf16 = _projector_linear_fp32_then_bf16( - torch, linear_input, linear.weight.detach(), linear.bias.detach() - ) - fp32_interpretation_is_distinct = not torch.equal( - fp32_then_bf16, projected_output - ) - if not fp32_interpretation_is_distinct: - raise StarVLAError( - "official PI_v3 projector sample does not distinguish BF16 operands from " - "FP32 Linear followed by BF16 output rounding" - ) - captures["projector_autocast_contract"] = dict(PROJECTOR_AUTOCAST_CONTRACT) - - captures["raw_qwen_taps"] = [ - raw_layer_outputs[index] for index in range(EXPECTED_LAYER_COUNT) - ] - captures["result_norm_diagnostic"] = final_norm["value"] - if captures["timestep_ids"] != EXPECTED_TIMESTEP_IDS: - raise StarVLAError( - f"official PI_v3 timestep order changed: {captures['timestep_ids']}" - ) - for name in ( - "conditioning_taps", - "raw_qwen_taps", - "result_norm_diagnostic", - "projected_hidden_taps", - "initial_noise", - ): - values = captures[name] if isinstance(captures[name], list) else [captures[name]] - if any(value.dtype != torch.bfloat16 for value in values): - raise StarVLAError(f"official PI_v3 {name} boundary is no longer BF16") - expected_noise_shape = (1, EXPECTED_ACTION_HORIZON, EXPECTED_ACTION_DIM) - if tuple(captures["initial_noise"].shape) != expected_noise_shape: - raise StarVLAError( - f"official PI_v3 initial noise shape mismatch: {tuple(captures['initial_noise'].shape)}" - ) - if captures["raw_policy"].dtype != torch.float32: - raise StarVLAError(f"official PI_v3 policy output is not FP32: {captures['raw_policy'].dtype}") - normalized = np.asarray(result.get("normalized_actions")) - raw_policy, _ = _tensor_to_array(captures["raw_policy"]) - if normalized.shape != expected_noise_shape or not np.array_equal(normalized, raw_policy): - raise StarVLAError("official normalized_actions differ from the captured PI_v3 policy output") - if not np.isfinite(normalized).all(): - raise StarVLAError("official PI_v3 policy produced NaN or infinite actions") - captures["normalized_actions"] = np.ascontiguousarray(normalized, dtype=np.float32) - return captures - - -def _image_pixel_sha256(image: Any) -> str: - header = _canonical_json({"mode": image.mode, "size": list(image.size)}) - return _sha256_bytes(header + b"\x00" + image.tobytes()) - - -def _image_record(path: Path, image: Any) -> dict[str, Any]: - return { - "source_path": str(path.resolve()), - "source_size": path.stat().st_size, - "source_sha256": sha256_file(path), - "decoded_mode": image.mode, - "decoded_size": list(image.size), - "decoded_pixel_sha256": _image_pixel_sha256(image), - } - - -def _processed_image_records(images: Sequence[Any]) -> list[dict[str, Any]]: - return [ - { - "index": index, - "mode": image.mode, - "size": list(image.size), - "pixel_sha256": _image_pixel_sha256(image), - } - for index, image in enumerate(images) - ] - - -def _render_model_prompt(framework: Any, processed_images: Sequence[Any], instruction: str) -> str: - messages = [ - { - "role": "user", - "content": [ - *({"type": "image", "image": image} for image in processed_images), - {"type": "text", "text": instruction}, - ], - } - ] - rendered = framework.qwen_vl_interface.processor.apply_chat_template( - messages, - tokenize=False, - add_generation_prompt=True, - ) - if not isinstance(rendered, str): - raise StarVLAError(f"official processor returned a non-string prompt: {type(rendered)}") - return rendered - - -def _processor_patch_contract(framework: Any) -> dict[str, int]: - vision_config = framework.qwen_vl_interface.model.config.vision_config - image_processor = framework.qwen_vl_interface.processor.image_processor - - def positive_int(value: Any, label: str) -> int: - if isinstance(value, bool) or not isinstance(value, int) or value <= 0: - raise StarVLAError(f"official Qwen processor {label} is not a positive integer: {value!r}") - return value - - channel_count = positive_int(vision_config.in_channels, "vision in_channels") - vision_patch_size = positive_int(vision_config.patch_size, "vision patch_size") - vision_temporal_patch_size = positive_int( - vision_config.temporal_patch_size, - "vision temporal_patch_size", - ) - processor_patch_size = positive_int(image_processor.patch_size, "image processor patch_size") - processor_temporal_patch_size = positive_int( - image_processor.temporal_patch_size, - "image processor temporal_patch_size", - ) - if ( - processor_patch_size != vision_patch_size - or processor_temporal_patch_size != vision_temporal_patch_size - ): - raise StarVLAError("official Qwen vision and image-processor patch contracts disagree") - contract = { - "channel_count": channel_count, - "patch_size": vision_patch_size, - "temporal_patch_size": vision_temporal_patch_size, - "pixel_patch_width": ( - channel_count * vision_temporal_patch_size * vision_patch_size * vision_patch_size - ), - } - expected = { - "channel_count": 3, - "patch_size": 16, - "temporal_patch_size": 2, - "pixel_patch_width": 1536, - } - if contract != expected: - raise StarVLAError(f"official Qwen processor patch contract changed: {contract}") - return contract - - -def _stack_taps(values: Sequence[Any], *, expected_width: int, label: str) -> Any: - import torch - - if len(values) != EXPECTED_LAYER_COUNT: - raise StarVLAError(f"{label} must contain exactly 36 tensors") - first_shape = tuple(values[0].shape) - if len(first_shape) != 3 or first_shape[0] != 1 or first_shape[2] != expected_width: - raise StarVLAError(f"unexpected {label} tensor shape: {first_shape}") - if any(tuple(value.shape) != first_shape for value in values): - raise StarVLAError(f"{label} tensors do not share one shape") - return torch.stack(list(values), dim=0).squeeze(1) - - -def _canonicalize_qwen_discrete_inputs( - qwen_inputs: Mapping[str, Any], -) -> tuple[np.ndarray, np.ndarray, np.ndarray, dict[str, str]]: - input_ids, input_dtype = _tensor_to_array(qwen_inputs["input_ids"]) - attention_mask, mask_dtype = _tensor_to_array(qwen_inputs["attention_mask"]) - image_grid, grid_dtype = _tensor_to_array(qwen_inputs["image_grid_thw"]) - source_dtypes = { - "input_ids": input_dtype, - "attention_mask": mask_dtype, - "image_grid_thw": grid_dtype, - } - drifted = [name for name, dtype in source_dtypes.items() if dtype != "int64"] - if drifted: - raise StarVLAError( - "official Qwen discrete processor inputs must originate as torch.int64: " - + ", ".join(f"{name}={source_dtypes[name]}" for name in drifted) - ) - if input_ids.ndim != 2 or input_ids.shape[0] != 1 or attention_mask.shape != input_ids.shape: - raise StarVLAError("official Qwen token/mask shape changed") - if not np.all((attention_mask == 0) | (attention_mask == 1)): - raise StarVLAError("official Qwen attention_mask contains values outside {0, 1}") - if not np.all(attention_mask == 1): - raise StarVLAError("single-sample official PI_v3 attention_mask must keep every token") - if image_grid.ndim != 2 or image_grid.shape != (1, 3): - raise StarVLAError(f"official Qwen image_grid_thw shape changed: {image_grid.shape}") - if np.any(image_grid <= 0): - raise StarVLAError( - f"official Qwen image_grid_thw contains non-positive values: {image_grid.tolist()}" - ) - return input_ids[0], attention_mask[0].astype(np.bool_), image_grid, source_dtypes - - -def _validate_qwen_pixel_values( - pixel_values: np.ndarray, - *, - source_dtype: Any, - image_grid: np.ndarray, - pixel_patch_width: int, -) -> None: - expected_patch_count = math.prod(int(value) for value in image_grid.flat) - if ( - pixel_values.dtype != np.float32 - or pixel_values.ndim != 2 - or pixel_values.shape[0] != expected_patch_count - or pixel_values.shape[1] != pixel_patch_width - or source_dtype != "float32" - ): - raise StarVLAError( - "official processor pixel_values are not source-FP32 patches matching " - f"image_grid_thw and width {pixel_patch_width}" - ) - - -def _build_arrays( - captures: Mapping[str, Any], - unnormalized: np.ndarray, - *, - pixel_patch_width: int, -) -> tuple[dict[str, np.ndarray], dict[str, Any]]: - arrays: dict[str, np.ndarray] = {} - records: dict[str, Any] = {} - - def add(name: str, value: Any, *, source_dtype: str | None = None) -> None: - if isinstance(value, np.ndarray): - array = np.ascontiguousarray(value) - inferred_dtype = None - else: - array, inferred_dtype = _tensor_to_array(value) - arrays[name] = array - records[name] = _array_record(array, source_dtype=source_dtype or inferred_dtype) - - qwen_inputs = captures["qwen_inputs"] - required_qwen_inputs = {"input_ids", "attention_mask", "image_grid_thw", "pixel_values"} - missing_qwen_inputs = sorted(required_qwen_inputs - set(qwen_inputs)) - if missing_qwen_inputs: - raise StarVLAError( - f"official Qwen preprocessing did not produce required tensors: {missing_qwen_inputs}" - ) - input_ids, attention_mask, image_grid, source_dtypes = _canonicalize_qwen_discrete_inputs( - qwen_inputs - ) - add("input_ids", input_ids, source_dtype=source_dtypes["input_ids"]) - add("attention_mask", attention_mask, source_dtype=source_dtypes["attention_mask"]) - add("image_grid_thw", image_grid, source_dtype=source_dtypes["image_grid_thw"]) - for key, tensor in sorted(qwen_inputs.items()): - if key not in {"input_ids", "attention_mask", "image_grid_thw"}: - add(f"qwen_input__{key}", tensor) - pixel_values = arrays["qwen_input__pixel_values"] - _validate_qwen_pixel_values( - pixel_values, - source_dtype=records["qwen_input__pixel_values"].get("source_dtype"), - image_grid=image_grid, - pixel_patch_width=pixel_patch_width, - ) - add( - "conditioning_taps", - _stack_taps(captures["conditioning_taps"], expected_width=EXPECTED_QWEN_HIDDEN_DIM, label="conditioning taps"), - ) - add( - "raw_qwen_taps", - _stack_taps(captures["raw_qwen_taps"], expected_width=EXPECTED_QWEN_HIDDEN_DIM, label="raw Qwen taps"), - ) - result_norm, result_norm_dtype = _tensor_to_array(captures["result_norm_diagnostic"]) - if result_norm.shape != (1, input_ids.shape[0], EXPECTED_QWEN_HIDDEN_DIM): - raise StarVLAError( - f"official Qwen result_norm diagnostic shape changed: {result_norm.shape}" - ) - add( - "result_norm_diagnostic", - np.ascontiguousarray(result_norm[0]), - source_dtype=result_norm_dtype, - ) - add( - "projected_hidden_taps", - _stack_taps( - captures["projected_hidden_taps"], - expected_width=EXPECTED_PROJECTED_HIDDEN_DIM, - label="projected hidden taps", - ), - ) - add("initial_noise", captures["initial_noise"]) - add("normalized_actions", captures["normalized_actions"]) - add("unnormalized_actions", np.ascontiguousarray(unnormalized, dtype=np.float32)) - return arrays, records - - -def _runtime_record(torch: Any, transformers: Any, device: str, recorder_probe: Mapping[str, Any]) -> dict[str, Any]: - cuda_device = torch.device(device) - index = cuda_device.index if cuda_device.index is not None else torch.cuda.current_device() - properties = torch.cuda.get_device_properties(index) - return { - "python": platform.python_version(), - "platform": platform.platform(), - "torch": torch.__version__, - "torchvision": _distribution_version("torchvision"), - "transformers": transformers.__version__, - "numpy": np.__version__, - "diffusers": _distribution_version("diffusers"), - "tokenizers": _distribution_version("tokenizers"), - "pillow": _distribution_version("Pillow"), - "omegaconf": _distribution_version("omegaconf"), - "accelerate": _distribution_version("accelerate"), - "safetensors": _distribution_version("safetensors"), - "official_environment_freeze": dict(OFFICIAL_ENVIRONMENT_FREEZE), - "cuda_runtime": torch.version.cuda, - "cudnn": torch.backends.cudnn.version(), - "device": str(cuda_device), - "device_name": properties.name, - "compute_capability": [properties.major, properties.minor], - "qwen3vl_recorder_probe": dict(recorder_probe), - } - - -def _copy_inputs(staging: Path, image_paths: Sequence[Path]) -> list[str]: - inputs_dir = staging / "inputs" - inputs_dir.mkdir() - relative_paths = [] - for index, source in enumerate(image_paths): - suffix = source.suffix.lower() if source.suffix else ".img" - destination = inputs_dir / f"image-{index:02d}{suffix}" - shutil.copyfile(source, destination) - relative_paths.append(destination.relative_to(staging).as_posix()) - return relative_paths - - -def _source_asset_hashes(entry: Mapping[str, Any], *, staged: bool = False) -> dict[str, str]: - overrides = entry.get("staged_overrides", {}) if staged else {} - return { - relative: overrides.get(relative, record)["sha256"] - for relative, record in entry["file_hashes"].items() - } - - -def write_golden( - *, - output_dir: Path, - paths: Mapping[str, Any], - framework: Any, - config: Mapping[str, Any], - recorder_probe: Mapping[str, Any], - image_paths: Sequence[Path], - source_image_records: Sequence[Mapping[str, Any]], - task: str, - unnorm_key: str, - captures: Mapping[str, Any], - unnormalized: np.ndarray, -) -> Path: - import torch - import transformers - - output_dir = output_dir.resolve() - if output_dir.exists(): - raise StarVLAError(f"golden output directory already exists: {output_dir}") - output_dir.parent.mkdir(parents=True, exist_ok=True) - processor_patch_contract = _processor_patch_contract(framework) - arrays, array_records = _build_arrays( - captures, - unnormalized, - pixel_patch_width=processor_patch_contract["pixel_patch_width"], - ) - input_ids = arrays["input_ids"] - attention_mask = arrays["attention_mask"] - if captures["framework_instructions"] != [task]: - raise StarVLAError( - f"official PI_v3 framework instruction changed: {captures['framework_instructions']!r}" - ) - model_instruction = expected_model_instruction(config, task) - rendered_prompt = _render_model_prompt(framework, captures["processed_images"], model_instruction) - token_strings = framework.qwen_vl_interface.processor.tokenizer.convert_ids_to_tokens(input_ids.tolist()) - runtime_contract = expected_runtime_contract() - runtime_contract_sha = _sha256_bytes(_canonical_json(runtime_contract)) - identity = { - "schema_version": GOLDEN_SCHEMA_VERSION, - "variant": SUPPORTED_VARIANT, - "checkpoint_sha256": paths["variant"]["checkpoint"]["sha256"], - "starvla_revision": paths["catalog"]["source_revisions"]["starvla"], - "qwen_revision": paths["qwen"]["revision"], - "runtime_contract_sha256": runtime_contract_sha, - "task": task, - "unnorm_key": unnorm_key, - "seed": SEED, - "images": [record["source_sha256"] for record in source_image_records], - } - golden_id = _sha256_bytes(_canonical_json(identity)) - - with tempfile.TemporaryDirectory(prefix=f".{output_dir.name}.", dir=output_dir.parent) as temporary: - staging = Path(temporary) - copied_images = _copy_inputs(staging, image_paths) - tensor_path = staging / "tensors.npz" - np.savez(tensor_path, **arrays) - image_records = [] - for index, record in enumerate(source_image_records): - copied = staging / copied_images[index] - image_records.append( - { - **record, - "artifact": copied_images[index], - "artifact_size": copied.stat().st_size, - "artifact_sha256": sha256_file(copied), - } - ) - - variant = paths["variant"] - qwen = paths["qwen"] - manifest: dict[str, Any] = { - "schema_version": GOLDEN_SCHEMA_VERSION, - "kind": GOLDEN_KIND, - "golden_id": golden_id, - "created_utc": dt.datetime.now(dt.timezone.utc).isoformat(), - "variant": SUPPORTED_VARIANT, - "model_type": variant["model_type"], - "source": { - "catalog": str(paths["catalog_path"]), - "catalog_sha256": sha256_file(paths["catalog_path"]), - "bundle_uuid": official_bundle_uuid(variant, paths["catalog"]), - "starvla_repo_revision": paths["catalog"]["source_revisions"]["starvla"], - "starvla_checkout": str(paths["source_dir"]), - "checkpoint_repo_id": variant["repo_id"], - "checkpoint_revision": variant["revision"], - "checkpoint_path": str(paths["checkpoint"]), - "checkpoint_size": variant["checkpoint"]["size"], - "checkpoint_sha256": variant["checkpoint"]["sha256"], - "policy_assets": _source_asset_hashes(variant), - "qwen_repo_id": qwen["repo_id"], - "qwen_revision": qwen["revision"], - "qwen_runtime_assets": _source_asset_hashes(qwen), - "qwen_converted_component_assets": _source_asset_hashes(qwen, staged=True), - }, - "runtime": _runtime_record( - torch, transformers, str(next(framework.parameters()).device), recorder_probe - ), - "determinism": { - "seed": SEED, - "rng_reset_immediately_before_predict": True, - "initial_noise_saved_explicitly": True, - "torch_deterministic_algorithms": True, - "cublas_workspace_config": os.environ.get("CUBLAS_WORKSPACE_CONFIG"), - "cuda_matmul_allow_tf32": False, - "cuda_matmul_allow_bf16_reduced_precision_reduction": False, - "cudnn_allow_tf32": False, - "cudnn_benchmark": False, - "attention_implementation": "sdpa", - }, - "compatibility": { - "qwen_bootstrap": ( - "config-only topology construction; all persistent parameters are then populated by " - "strict loading of the pinned official checkpoint" - ), - "effective_config": "config.yaml with pinned checkpoint-derived PI_v3 compatibility fixes", - "projector_autocast": captures["projector_autocast_contract"], - "policy_boundary_casts": { - "projected_hidden": "BF16 output widened exactly to FP32 at DiT cross-attention input", - "initial_noise": "BF16 torch.randn output widened exactly to FP32 at action encoder input", - "reason": ( - "the released source requests CUDA autocast(dtype=float32), which PyTorch 2.6 disables; " - "these two explicit boundary casts realize its declared FP32 policy path" - ), - }, - }, - "input": { - "task": task, - "unnorm_key": unnorm_key, - "state": None, - "images": image_records, - "processed_images": _processed_image_records(captures["processed_images"]), - }, - "model_contract": { - "framework_class": f"{type(framework).__module__}.{type(framework).__name__}", - "action_horizon": EXPECTED_ACTION_HORIZON, - "action_dim": EXPECTED_ACTION_DIM, - "qwen_hidden_dim": EXPECTED_QWEN_HIDDEN_DIM, - "qwen_layer_count": EXPECTED_LAYER_COUNT, - "projected_hidden_dim": EXPECTED_PROJECTED_HIDDEN_DIM, - "hidden_tuple_indices": list(range(1, 37)), - "conditioning_tap_names": list(CONDITIONING_TAP_NAMES), - "raw_tap_names": list(RAW_TAP_NAMES), - "diagnostic_tap_names": [FINAL_NORM_DIAGNOSTIC_NAME], - "tap_layout": "layer_token_hidden", - "conditioning_semantics": CONDITIONING_SEMANTICS, - "result_norm_role": "golden_only_diagnostic_not_conditioning_or_candidate_gate", - "timestep_ids": EXPECTED_TIMESTEP_IDS, - "initial_noise_dtype": "bfloat16", - "policy_compute_dtype": "float32", - "state_input_active": False, - "runtime_contract": runtime_contract, - "runtime_contract_sha256": runtime_contract_sha, - }, - "prompt": { - "framework_instruction": task, - "model_instruction": model_instruction, - "rendered_chat_template": rendered_prompt, - "action_token_mode": "none", - }, - "processor": { - "image_grid_thw": arrays["image_grid_thw"].tolist(), - "pixel_values_shape": list(arrays["qwen_input__pixel_values"].shape), - "patch_contract": processor_patch_contract, - "qwen_input_array_names": sorted( - name for name in arrays if name.startswith("qwen_input__") - ), - "smart_resize_values_are_observed_not_assumed": True, - }, - "tokens": { - "input_ids": input_ids.tolist(), - "attention_mask": attention_mask.tolist(), - "token_strings": token_strings, - }, - "outputs": { - "normalized_actions": arrays["normalized_actions"].tolist(), - "unnormalized_actions": arrays["unnormalized_actions"].tolist(), - }, - "artifacts": { - "tensors": { - "path": tensor_path.name, - "size": tensor_path.stat().st_size, - "sha256": sha256_file(tensor_path), - "encoding": "numpy_npz_stored", - "arrays": array_records, - } - }, - } - manifest["integrity"] = { - "canonicalization": "utf8_json_sort_keys_compact_excluding_integrity", - "manifest_payload_sha256": _sha256_bytes(_canonical_json(manifest)), - } - manifest_path = staging / "golden.json" - manifest_path.write_text( - json.dumps(manifest, ensure_ascii=False, indent=2, sort_keys=True) + "\n", - encoding="utf-8", - ) - Path(temporary).replace(output_dir) - return output_dir / "golden.json" - - -def build_parser() -> argparse.ArgumentParser: - parser = argparse.ArgumentParser( - description="Generate an auditable golden from the pinned official StarVLA Qwen3-VL PI_v3 checkpoint." - ) - parser.add_argument("--image", action="append", default=[], type=Path, help="The single 224x224 RGB image") - parser.add_argument("--task", help="Robot task instruction") - parser.add_argument("--unnorm-key", choices=("oxe_bridge", "oxe_rt1")) - parser.add_argument("--output-dir", type=Path) - parser.add_argument("--checkpoint-root", type=Path, default=Path("ckpts/starvla")) - parser.add_argument("--starvla-source", type=Path, default=None) - parser.add_argument("--device", default="cuda:0") - parser.add_argument( - "--preflight-only", - action="store_true", - help="Verify pinned source/assets/checkpoint/runtime/recorder semantics without allocating the full model", - ) - return parser - - -def _load_images(image_paths: Iterable[Path]) -> tuple[list[Any], list[dict[str, Any]]]: - try: - from PIL import Image - except ImportError as exc: - raise StarVLAError("Pillow is required to load oracle images") from exc - images = [] - records = [] - for path in image_paths: - path = path.resolve() - _ensure_regular_file(path, label="input image") - try: - with Image.open(path) as opened: - opened.load() - image = opened.copy() - except (OSError, ValueError) as exc: - raise StarVLAError(f"failed to decode input image {path}: {exc}") from exc - if image.mode != "RGB": - raise StarVLAError(f"official PI_v3 golden input must already be RGB, got mode {image.mode!r}") - if image.size != (224, 224): - raise StarVLAError(f"official PI_v3 golden input must be exactly 224x224, got {image.size}") - images.append(image) - records.append(_image_record(path, image)) - return images, records - - -def main(argv: Sequence[str] | None = None) -> int: - args = build_parser().parse_args(argv) - _require_isolated_python() - if not args.preflight_only: - missing = [ - name - for name, value in ( - ("--image", args.image), - ("--task", args.task), - ("--unnorm-key", args.unnorm_key), - ("--output-dir", args.output_dir), - ) - if not value - ] - if missing: - raise StarVLAError("golden generation requires " + ", ".join(missing)) - if len(args.image) != 1: - raise StarVLAError("the released PI_v3 checkpoint requires exactly one image") - if not args.task.strip(): - raise StarVLAError("--task must not be empty") - checkpoint_root = args.checkpoint_root.resolve() - output_dir = args.output_dir.resolve() - if output_dir == checkpoint_root or checkpoint_root in output_dir.parents: - raise StarVLAError("--output-dir must not be inside the pinned checkpoint source tree") - - checkpoint_root = args.checkpoint_root.resolve() - source_dir = args.starvla_source or checkpoint_root / "source" / "starvla" - paths = validate_official_inputs( - checkpoint_root=checkpoint_root, - source_dir=source_dir, - catalog_path=DEFAULT_CATALOG, - ) - try: - import torch - import transformers - except ImportError as exc: - raise StarVLAError(f"official StarVLA runtime dependency is missing: {exc}") from exc - validate_runtime_versions( - torch_version=torch.__version__, - torchvision_version=_distribution_version("torchvision"), - transformers_version=transformers.__version__, - numpy_version=np.__version__, - diffusers_version=_distribution_version("diffusers"), - tokenizers_version=_distribution_version("tokenizers"), - pillow_version=_distribution_version("Pillow"), - omegaconf_version=_distribution_version("omegaconf"), - accelerate_version=_distribution_version("accelerate"), - safetensors_version=_distribution_version("safetensors"), - ) - _configure_determinism(torch, seed=SEED, device=args.device) - recorder_probe = verify_transformers_qwen3vl_recorder_semantics(torch, transformers) - expected_runtime_contract() - if args.preflight_only: - verify_official_framework_import(paths) - print("Pinned StarVLA PI_v3 oracle preflight passed.") - return 0 - - images, source_image_records = _load_images(args.image) - framework, config = load_official_framework(paths, device=args.device) - captures = run_official_forward(framework, images=images, task=args.task, seed=SEED) - - source_dir = Path(paths["source_dir"]) - sys.path.insert(0, str(source_dir)) - try: - from deployment.model_server import policy_norm_processor - - _assert_module_origin(policy_norm_processor, source_dir) - normalizer = policy_norm_processor.PolicyNormProcessor( - str(paths["checkpoint"]), unnorm_key=args.unnorm_key - ) - normalized = captures["normalized_actions"] - unnormalized = np.asarray(normalizer.unapply_actions(normalized[0]))[None, ...] - if unnormalized.shape != normalized.shape or not np.isfinite(unnormalized).all(): - raise StarVLAError( - f"official action unnormalization returned invalid values/shape: {unnormalized.shape}" - ) - finally: - if sys.path and sys.path[0] == str(source_dir): - del sys.path[0] - - manifest = write_golden( - output_dir=args.output_dir, - paths=paths, - framework=framework, - config=config, - recorder_probe=recorder_probe, - image_paths=args.image, - source_image_records=source_image_records, - task=args.task, - unnorm_key=args.unnorm_key, - captures=captures, - unnormalized=np.ascontiguousarray(unnormalized, dtype=np.float32), - ) - print(f"Wrote official StarVLA PI_v3 golden: {manifest}") - return 0 - - -if __name__ == "__main__": - try: - raise SystemExit(main()) - except StarVLAError as exc: - raise SystemExit(f"error: {exc}") from exc diff --git a/tools/hf2gguf/starvla/generate_starvla_qwen25_fast_golden.py b/tools/hf2gguf/starvla/generate_starvla_qwen25_fast_golden.py deleted file mode 100644 index 4bdd889..0000000 --- a/tools/hf2gguf/starvla/generate_starvla_qwen25_fast_golden.py +++ /dev/null @@ -1,765 +0,0 @@ -#!/usr/bin/env python3 -"""Generate a CUDA local-Python action golden from the official Qwen2.5 FAST .pt.""" - -from __future__ import annotations - -import argparse -import hashlib -import importlib.metadata -import json -import math -import os -import platform -import random -import sys -from pathlib import Path -from typing import Any, Mapping, Sequence - - -TOOLS_DIR = Path(__file__).resolve().parent -if str(TOOLS_DIR) not in sys.path: - sys.path.insert(0, str(TOOLS_DIR)) - -from convert_starvla_qwen25_fast import ( # noqa: E402 - ACTION_DIM, - ACTION_HORIZON, - ACTION_TOKEN_MAX, - ACTION_TOKEN_MIN, - BACKBONE, - COT_PROMPT, - FAST_CODEC_ASSET_KEY, - FRAMEWORK, - GENERATION_CONTRACT, - MODEL_TYPE, - STAGING_MANIFEST_FILENAME, - VARIANT_KEY, - validate_fast_codec, - validate_staging_manifest, -) -from starvla_checkpoint import ( # noqa: E402 - DEFAULT_CATALOG, - StarVLAError, - atomic_write_json, - get_variant, - load_catalog, - official_bundle_uuid, - sha256_file, - verify_checkpoint_file, -) - - -SCHEMA_VERSION = 2 -GOLDEN_KIND = "starvla_qwen25_fast_local_python_action_golden" -DEFAULT_SEED = 42 -UNNORM_KEYS = ("bridge_dataset", "fractal20220817_data") -EXPECTED_RUNTIME_VERSIONS = { - "torch": "2.6.0", - "torchvision": "0.21.0", - "transformers": "4.57.0", - "numpy": "1.26.4", - "qwen-vl-utils": "0.0.14", -} - - -def canonical_sha256(value: Any) -> str: - payload = json.dumps(value, sort_keys=True, separators=(",", ":")).encode("utf-8") - return hashlib.sha256(payload).hexdigest() - - -def valid_sha256(value: Any) -> bool: - return ( - isinstance(value, str) - and len(value) == 64 - and all(character in "0123456789abcdef" for character in value) - ) - - -def validate_runtime_versions(actual: Mapping[str, str]) -> None: - mismatches = [] - for name, expected in EXPECTED_RUNTIME_VERSIONS.items(): - version = str(actual.get(name, "missing")).split("+", 1)[0] - if version != expected: - mismatches.append(f"{name}: expected {expected}, got {version}") - if mismatches: - raise StarVLAError( - "Qwen2.5 FAST local-Python runtime version mismatch: " - + "; ".join(mismatches) - ) - - -def distribution_version(name: str) -> str: - try: - return importlib.metadata.version(name) - except importlib.metadata.PackageNotFoundError: - return "missing" - - -def render_prompt(instruction: str) -> str: - if not isinstance(instruction, str) or not instruction.strip(): - raise StarVLAError("FAST instruction must be a non-empty string") - return COT_PROMPT.replace("{instruction}", instruction) - - -def build_messages(image: Any, instruction: str) -> list[dict[str, Any]]: - return [ - { - "role": "user", - "content": [ - {"type": "image", "image": image}, - {"type": "text", "text": render_prompt(instruction)}, - ], - } - ] - - -def extract_action_token_ids( - generated_ids: Sequence[Sequence[int]], -) -> list[list[int]]: - result = [] - for row in generated_ids: - tokens = [] - for value in row: - if not isinstance(value, int) or isinstance(value, bool): - raise StarVLAError("generated token IDs must be integers") - if ACTION_TOKEN_MIN <= value <= ACTION_TOKEN_MAX: - tokens.append(value) - result.append(tokens) - return result - - -def map_vlm_to_fast_ids( - batch_action_token_ids: Sequence[Sequence[int]], -) -> list[list[int]]: - result = [] - for row in batch_action_token_ids: - fast_ids = [token_id - ACTION_TOKEN_MIN for token_id in row] - if any(token_id < 0 or token_id > 2047 for token_id in fast_ids): - raise StarVLAError("generated action token is outside the FAST vocabulary") - result.append(fast_ids) - return result - - -def validate_actions(value: Any, *, name: str) -> list[list[list[float]]]: - try: - import numpy as np - except ImportError as exc: - raise StarVLAError("NumPy is required for FAST action validation") from exc - actions = np.asarray(value, dtype=np.float64) - expected_shape = (1, ACTION_HORIZON, ACTION_DIM) - if actions.shape != expected_shape: - raise StarVLAError( - f"{name} has shape {list(actions.shape)}, expected {list(expected_shape)}" - ) - if not np.isfinite(actions).all(): - raise StarVLAError(f"{name} contains a non-finite value") - return actions.tolist() - - -def validate_normalized_actions(value: Any) -> list[list[list[float]]]: - return validate_actions(value, name="FAST normalized actions") - - -def load_normalization_profile( - dataset_statistics: Path, - unnorm_key: str, -) -> dict[str, Any]: - if unnorm_key not in UNNORM_KEYS: - raise StarVLAError( - f"FAST --unnorm-key must be one of {list(UNNORM_KEYS)}, got {unnorm_key!r}" - ) - try: - statistics = json.loads(dataset_statistics.read_text(encoding="utf-8")) - except (OSError, UnicodeError, json.JSONDecodeError) as exc: - raise StarVLAError( - f"failed to load FAST dataset statistics {dataset_statistics}: {exc}" - ) from exc - if not isinstance(statistics, dict) or set(statistics) != set(UNNORM_KEYS): - raise StarVLAError("FAST dataset statistics profile set is incompatible") - try: - action = statistics[unnorm_key]["action"] - q01 = [float(value) for value in action["q01"]] - q99 = [float(value) for value in action["q99"]] - mask = list(action["mask"]) - except (KeyError, TypeError, ValueError) as exc: - raise StarVLAError("FAST action normalization statistics are malformed") from exc - if ( - len(q01) != ACTION_DIM - or len(q99) != ACTION_DIM - or mask != [True] * 6 + [False] - or not all(math.isfinite(value) for value in q01 + q99) - or any(high < low for low, high in zip(q01[:6], q99[:6])) - ): - raise StarVLAError("FAST action normalization profile is incompatible") - return { - "profile": unnorm_key, - "action_q01": q01, - "action_q99": q99, - "action_mask": mask, - "continuous_dimensions": [0, 1, 2, 3, 4, 5], - "binary_dimensions": [6], - "binary_threshold": 0.5, - "binary_comparison": "gt", - "clip_actions": False, - } - - -def unnormalize_actions( - normalized_actions: Any, - normalization: Mapping[str, Any], -) -> list[list[list[float]]]: - try: - import numpy as np - except ImportError as exc: - raise StarVLAError("NumPy is required for FAST action unnormalization") from exc - normalized = np.asarray( - validate_actions(normalized_actions, name="FAST normalized actions"), - dtype=np.float32, - ) - q01 = np.asarray(normalization["action_q01"], dtype=np.float32) - q99 = np.asarray(normalization["action_q99"], dtype=np.float32) - result = np.empty_like(normalized, dtype=np.float32) - result[..., :6] = (normalized[..., :6] + np.float32(1.0)) * np.float32( - 0.5 - ) * (q99[:6] - q01[:6]) + q01[:6] - result[..., 6] = (normalized[..., 6] > np.float32(0.5)).astype(np.float32) - return validate_actions(result.tolist(), name="FAST unnormalized actions") - - -def validate_fast_token_rows( - fast_processor: Any, - batch_fast_ids: Sequence[Sequence[int]], -) -> None: - expected_coefficients = ACTION_HORIZON * ACTION_DIM - for index, row in enumerate(batch_fast_ids): - try: - decoded = fast_processor.bpe_tokenizer.decode(list(row)) - except Exception as exc: - raise StarVLAError(f"FAST token row {index} cannot be decoded") from exc - if len(decoded) != expected_coefficients: - raise StarVLAError( - f"FAST token row {index} contains {len(decoded)} coefficients; " - f"expected {expected_coefficients}" - ) - - -def finalize_golden_id(value: dict[str, Any]) -> str: - payload = dict(value) - payload.pop("golden_id", None) - golden_id = canonical_sha256(payload) - value["golden_id"] = golden_id - return golden_id - - -def _require_regular_bound_file( - path_value: Any, - size_value: Any, - sha_value: Any, - *, - label: str, -) -> Path: - path = Path(str(path_value)) - if not path.is_absolute() or not path.is_file() or path.is_symlink(): - raise StarVLAError(f"{label} must be an absolute regular file") - if ( - not isinstance(size_value, int) - or isinstance(size_value, bool) - or path.stat().st_size != size_value - or not valid_sha256(sha_value) - or sha256_file(path) != sha_value - ): - raise StarVLAError(f"{label} no longer matches its bound size/SHA256") - return path - - -def validate_golden( - value: Any, - *, - verify_files: bool = False, - catalog_path: Path = DEFAULT_CATALOG, -) -> dict[str, Any]: - if not isinstance(value, dict) or value.get("kind") != GOLDEN_KIND: - raise StarVLAError("not a Qwen2.5 FAST local-Python golden") - catalog = load_catalog(catalog_path) - variant = get_variant(catalog, VARIANT_KEY) - qwen = catalog["shared_assets"][variant["qwen_asset"]] - codec = catalog["shared_assets"][FAST_CODEC_ASSET_KEY] - expected = { - "schema_version": SCHEMA_VERSION, - "variant": VARIANT_KEY, - "framework": FRAMEWORK, - "backbone": BACKBONE, - "model_type": MODEL_TYPE, - "bundle_uuid": official_bundle_uuid(variant, catalog), - "generation": GENERATION_CONTRACT, - } - mismatches = [ - f"{key}: expected {item!r}, got {value.get(key)!r}" - for key, item in expected.items() - if value.get(key) != item - ] - source = value.get("source") - input_record = value.get("input") - prompt = value.get("prompt") - normalization = value.get("normalization") - runtime = value.get("runtime") - result = value.get("result") - if not all( - isinstance(item, dict) - for item in ( - source, - input_record, - prompt, - normalization, - runtime, - result, - ) - ): - mismatches.append( - "source/input/prompt/normalization/runtime/result must be objects" - ) - if mismatches: - raise StarVLAError("invalid Qwen2.5 FAST golden: " + "; ".join(mismatches)) - - source_expected = { - "checkpoint_repo_id": variant["repo_id"], - "checkpoint_revision": variant["revision"], - "checkpoint_filename": Path(variant["checkpoint"]["path"]).name, - "checkpoint_size": variant["checkpoint"]["size"], - "checkpoint_sha256": variant["checkpoint"]["sha256"], - "qwen_repo_id": qwen["repo_id"], - "qwen_revision": qwen["revision"], - "fast_codec_repo_id": codec["repo_id"], - "fast_codec_revision": codec["revision"], - "fast_codec_files": { - relative: codec["file_hashes"][relative]["sha256"] - for relative in codec["files"] - }, - "starvla_revision": catalog["source_revisions"]["starvla"], - "llama_revision": catalog["source_revisions"]["llama_cpp"], - "weight_source": "official_original_pt_staged_exact_weights", - } - for key, expected_value in source_expected.items(): - if source.get(key) != expected_value: - mismatches.append(f"source.{key}: expected {expected_value!r}") - if source.get("bundle_uuid") != value["bundle_uuid"]: - mismatches.append("source.bundle_uuid") - if prompt != { - "chat_template_sha256": qwen["file_hashes"]["chat_template.jinja"][ - "sha256" - ] - }: - mismatches.append("prompt.chat_template_sha256") - - instruction = input_record.get("instruction") - if ( - not isinstance(instruction, str) - or input_record.get("framework_prompt") != render_prompt(instruction) - or input_record.get("unnorm_key") not in UNNORM_KEYS - or not valid_sha256(input_record.get("image_sha256")) - or input_record.get("prompt_length", 0) <= 0 - or canonical_sha256(input_record.get("input_ids")) - != input_record.get("input_ids_sha256") - ): - mismatches.append("input prompt/image/profile/token contract") - - expected_profile = load_normalization_profile( - Path(source["dataset_statistics_path"]), - input_record["unnorm_key"], - ) - for key, expected_value in expected_profile.items(): - if normalization.get(key) != expected_value: - mismatches.append(f"normalization.{key}") - if ( - source.get("dataset_statistics_sha256") - != variant["file_hashes"]["dataset_statistics.json"]["sha256"] - or normalization.get("source_sha256") - != source.get("dataset_statistics_sha256") - ): - mismatches.append("normalization source SHA256") - - validate_runtime_versions(runtime) - if ( - runtime.get("backend") != "cuda" - or runtime.get("full_gpu_model") is not True - or runtime.get("dtype") != "bfloat16" - or runtime.get("attn_implementation") != "sdpa" - or runtime.get("tf32") is not False - ): - mismatches.append("runtime CUDA/dtype/attention contract") - - generated = result.get("generated_ids") - action_ids = result.get("action_token_ids") - fast_ids = result.get("fast_token_ids") - if not isinstance(generated, list) or len(generated) != 1: - mismatches.append("result.generated_ids") - elif extract_action_token_ids(generated) != action_ids: - mismatches.append("result.action_token_ids") - if isinstance(action_ids, list): - try: - if map_vlm_to_fast_ids(action_ids) != fast_ids: - mismatches.append("result.fast_token_ids") - except StarVLAError as exc: - mismatches.append(str(exc)) - if ( - not isinstance(action_ids, list) - or len(action_ids) != 1 - or not action_ids[0] - or not isinstance(fast_ids, list) - or len(fast_ids) != 1 - or not fast_ids[0] - ): - mismatches.append("result requires non-empty action/FAST token IDs") - for key in ("normalized_actions", "unnormalized_actions"): - try: - validate_actions(result.get(key), name=f"result.{key}") - except StarVLAError as exc: - mismatches.append(str(exc)) - if result.get("unnormalized_actions") != unnormalize_actions( - result.get("normalized_actions"), normalization - ): - mismatches.append("result.unnormalized_actions formula") - if ( - canonical_sha256(result.get("generated_ids")) - != result.get("generated_ids_sha256") - or canonical_sha256(result.get("normalized_actions")) - != result.get("normalized_actions_sha256") - or canonical_sha256(result.get("unnormalized_actions")) - != result.get("unnormalized_actions_sha256") - ): - mismatches.append("result canonical SHA256") - golden_id = value.get("golden_id") - payload = dict(value) - payload.pop("golden_id", None) - if not valid_sha256(golden_id) or canonical_sha256(payload) != golden_id: - mismatches.append("golden_id") - if mismatches: - raise StarVLAError("invalid Qwen2.5 FAST golden: " + "; ".join(mismatches)) - - if verify_files: - checkpoint = _require_regular_bound_file( - source["checkpoint_path"], - source["checkpoint_size"], - source["checkpoint_sha256"], - label="golden source checkpoint", - ) - verify_checkpoint_file(checkpoint, variant) - statistics_path = _require_regular_bound_file( - source["dataset_statistics_path"], - source["dataset_statistics_size"], - source["dataset_statistics_sha256"], - label="golden dataset statistics", - ) - if statistics_path.name != "dataset_statistics.json": - raise StarVLAError("golden dataset statistics filename is incompatible") - image_path = _require_regular_bound_file( - input_record["image_path"], - input_record["image_size"], - input_record["image_sha256"], - label="golden input image", - ) - if image_path.resolve() != Path(input_record["image_path"]): - raise StarVLAError("golden image path is not canonical") - manifest_path = _require_regular_bound_file( - source["staging_manifest_path"], - source["staging_manifest_size"], - source["staging_manifest_sha256"], - label="golden staging manifest", - ) - try: - manifest = json.loads(manifest_path.read_text(encoding="utf-8")) - except (OSError, json.JSONDecodeError) as exc: - raise StarVLAError(f"failed to parse golden staging manifest: {exc}") from exc - staging_dir = Path(source["staged_hf_path"]).parent - validate_staging_manifest(manifest, catalog, staging_dir) - if ( - Path(source["staged_hf_path"]) != (staging_dir / "hf").resolve() - or Path(manifest["source"]["checkpoint"]).resolve() != checkpoint.resolve() - or manifest["bundle_uuid"] != value["bundle_uuid"] - ): - raise StarVLAError("golden staged exact-weight binding is inconsistent") - codec_path = Path(str(source["fast_codec_path"])) - if ( - not codec_path.is_absolute() - or not codec_path.is_dir() - or codec_path.is_symlink() - ): - raise StarVLAError("golden FAST codec path is not a bound directory") - actual_codec = validate_fast_codec(codec_path, codec) - if actual_codec["files"] != source["fast_codec_files"]: - raise StarVLAError("golden FAST codec source binding changed") - return value - - -def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--checkpoint", type=Path, required=True) - parser.add_argument("--staged-hf", type=Path, required=True) - parser.add_argument("--staging-manifest", type=Path) - parser.add_argument("--fast-codec", type=Path, required=True) - parser.add_argument("--image", type=Path, required=True) - parser.add_argument("--instruction", required=True) - parser.add_argument("--unnorm-key", choices=UNNORM_KEYS, required=True) - parser.add_argument("--output", type=Path, required=True) - parser.add_argument("--catalog", type=Path, default=DEFAULT_CATALOG) - parser.add_argument( - "--attn-implementation", - choices=("sdpa",), - default="sdpa", - ) - parser.add_argument("--seed", type=int, default=DEFAULT_SEED) - return parser.parse_args(argv) - - -def main(argv: Sequence[str] | None = None) -> int: - args = parse_args(argv) - try: - if not sys.flags.isolated: - raise StarVLAError( - "Qwen2.5 FAST golden generation must run in isolated mode (`python -I`)" - ) - os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" - import numpy as np - import torch - import transformers - from PIL import Image - from qwen_vl_utils import process_vision_info - from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration - - if not torch.cuda.is_available(): - raise StarVLAError("Qwen2.5 FAST golden generation requires CUDA") - runtime_versions = { - "torch": torch.__version__, - "torchvision": distribution_version("torchvision"), - "transformers": transformers.__version__, - "numpy": np.__version__, - "qwen-vl-utils": distribution_version("qwen-vl-utils"), - } - validate_runtime_versions(runtime_versions) - catalog = load_catalog(args.catalog) - variant = get_variant(catalog, VARIANT_KEY) - qwen = catalog["shared_assets"][variant["qwen_asset"]] - codec_entry = catalog["shared_assets"][FAST_CODEC_ASSET_KEY] - verify_checkpoint_file(args.checkpoint, variant) - codec = validate_fast_codec(args.fast_codec.resolve(), codec_entry) - - manifest_path = ( - args.staging_manifest - or args.staged_hf.parent / STAGING_MANIFEST_FILENAME - ).resolve() - try: - manifest = json.loads(manifest_path.read_text(encoding="utf-8")) - except (OSError, json.JSONDecodeError) as exc: - raise StarVLAError( - f"failed to load staging manifest {manifest_path}: {exc}" - ) from exc - staging_dir = args.staged_hf.resolve().parent - validate_staging_manifest(manifest, catalog, staging_dir) - if ( - args.staged_hf.resolve() != staging_dir / "hf" - or Path(manifest["source"]["checkpoint"]).resolve() - != args.checkpoint.resolve() - ): - raise StarVLAError( - "staging manifest is bound to a different exact-weight source" - ) - dataset_statistics = args.checkpoint.resolve().parents[1] / "dataset_statistics.json" - normalization = load_normalization_profile( - dataset_statistics, args.unnorm_key - ) - normalization["source_sha256"] = sha256_file(dataset_statistics) - expected_stats_sha = variant["file_hashes"]["dataset_statistics.json"]["sha256"] - if normalization["source_sha256"] != expected_stats_sha: - raise StarVLAError("official FAST dataset statistics SHA256 changed") - if not args.image.is_file() or args.image.is_symlink(): - raise StarVLAError(f"missing FAST golden image: {args.image}") - if args.seed < 0: - raise StarVLAError("--seed must be non-negative") - - random.seed(args.seed) - np.random.seed(args.seed) - torch.manual_seed(args.seed) - torch.cuda.manual_seed_all(args.seed) - torch.use_deterministic_algorithms(True) - torch.backends.cuda.matmul.allow_tf32 = False - torch.backends.cudnn.allow_tf32 = False - torch.backends.cudnn.benchmark = False - - processor = AutoProcessor.from_pretrained( - args.staged_hf, local_files_only=True - ) - processor.tokenizer.padding_side = "left" - model = Qwen2_5_VLForConditionalGeneration.from_pretrained( - args.staged_hf, - local_files_only=True, - torch_dtype=torch.bfloat16, - attn_implementation=args.attn_implementation, - ).to("cuda") - model.eval() - - generation_file = json.loads( - (args.staged_hf / "generation_config.json").read_text(encoding="utf-8") - ) - for key, expected in GENERATION_CONTRACT.items(): - if key != "max_length" and generation_file.get(key) != expected: - raise StarVLAError( - f"staged generation_config drift at {key}: " - f"expected {expected!r}, got {generation_file.get(key)!r}" - ) - - with Image.open(args.image) as opened: - image = opened.convert("RGB") - messages = build_messages(image, args.instruction) - rendered_chat = processor.apply_chat_template( - messages, tokenize=False, add_generation_prompt=True - ) - image_inputs, video_inputs = process_vision_info([messages]) - inputs = processor( - text=[rendered_chat], - images=image_inputs, - videos=video_inputs, - padding=True, - return_tensors="pt", - ).to("cuda") - input_ids = inputs["input_ids"].detach().cpu().tolist() - prompt_length = int(inputs["input_ids"].shape[1]) - - with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16): - generated_tensor = model.generate( - **inputs, max_length=GENERATION_CONTRACT["max_length"] - ) - generated_ids = generated_tensor.detach().cpu().tolist() - action_token_ids = extract_action_token_ids(generated_ids) - fast_token_ids = map_vlm_to_fast_ids(action_token_ids) - - fast_processor = AutoProcessor.from_pretrained( - args.fast_codec, - trust_remote_code=True, - local_files_only=True, - ) - fast_processor.time_horizon = ACTION_HORIZON - fast_processor.action_dim = ACTION_DIM - decode_inputs = [row if row else None for row in fast_token_ids] - validate_fast_token_rows(fast_processor, fast_token_ids) - normalized_actions = validate_normalized_actions( - fast_processor.decode(decode_inputs) - ) - unnormalized = unnormalize_actions(normalized_actions, normalization) - - checkpoint = args.checkpoint.resolve() - image_path = args.image.resolve() - golden: dict[str, Any] = { - "schema_version": SCHEMA_VERSION, - "kind": GOLDEN_KIND, - "variant": VARIANT_KEY, - "framework": FRAMEWORK, - "backbone": BACKBONE, - "model_type": MODEL_TYPE, - "bundle_uuid": manifest["bundle_uuid"], - "source": { - "bundle_uuid": manifest["bundle_uuid"], - "checkpoint_repo_id": variant["repo_id"], - "checkpoint_revision": variant["revision"], - "checkpoint_filename": checkpoint.name, - "checkpoint_path": str(checkpoint), - "checkpoint_size": checkpoint.stat().st_size, - "checkpoint_sha256": variant["checkpoint"]["sha256"], - "dataset_statistics_path": str(dataset_statistics.resolve()), - "dataset_statistics_size": dataset_statistics.stat().st_size, - "dataset_statistics_sha256": expected_stats_sha, - "qwen_repo_id": qwen["repo_id"], - "qwen_revision": qwen["revision"], - "fast_codec_repo_id": codec["repo_id"], - "fast_codec_revision": codec["revision"], - "fast_codec_files": codec["files"], - "starvla_revision": catalog["source_revisions"]["starvla"], - "llama_revision": catalog["source_revisions"]["llama_cpp"], - "staged_hf_path": str(args.staged_hf.resolve()), - "staging_manifest_path": str(manifest_path), - "staging_manifest_size": manifest_path.stat().st_size, - "staging_manifest_sha256": sha256_file(manifest_path), - "fast_codec_path": str(args.fast_codec.resolve()), - "weight_source": "official_original_pt_staged_exact_weights", - }, - "prompt": { - "chat_template_sha256": qwen["file_hashes"][ - "chat_template.jinja" - ]["sha256"], - }, - "runtime": { - "python": platform.python_version(), - **runtime_versions, - "backend": "cuda", - "full_gpu_model": True, - "device": torch.cuda.get_device_name(torch.cuda.current_device()), - "cuda": torch.version.cuda, - "dtype": "bfloat16", - "attn_implementation": args.attn_implementation, - "tf32": False, - "torch_deterministic_algorithms": True, - "cublas_workspace_config": ":4096:8", - "seed": args.seed, - }, - "input": { - "image_path": str(image_path), - "image_size": image_path.stat().st_size, - "image_sha256": sha256_file(image_path), - "decoded_size": [image.width, image.height], - "instruction": args.instruction, - "unnorm_key": args.unnorm_key, - "framework_prompt": render_prompt(args.instruction), - "rendered_chat_template": rendered_chat, - "input_ids": input_ids, - "input_ids_sha256": canonical_sha256(input_ids), - "prompt_length": prompt_length, - }, - "normalization": normalization, - "generation": dict(GENERATION_CONTRACT), - "result": { - "generated_ids": generated_ids, - "continuation_ids": [ - row[prompt_length:] for row in generated_ids - ], - "generated_ids_sha256": canonical_sha256(generated_ids), - "action_token_ids": action_token_ids, - "fast_token_ids": fast_token_ids, - "normalized_actions": normalized_actions, - "normalized_actions_sha256": canonical_sha256(normalized_actions), - "unnormalized_actions": unnormalized, - "unnormalized_actions_sha256": canonical_sha256(unnormalized), - }, - } - finalize_golden_id(golden) - validate_golden( - golden, verify_files=True, catalog_path=args.catalog - ) - atomic_write_json(args.output, golden, overwrite=False) - print(f"Qwen2.5 FAST golden: {args.output}") - print( - json.dumps( - { - "golden_id": golden["golden_id"], - "prompt_tokens": prompt_length, - "generated_tokens": len(generated_ids[0]) - prompt_length, - "action_tokens": len(action_token_ids[0]), - "action_shape": [1, ACTION_HORIZON, ACTION_DIM], - "unnorm_key": args.unnorm_key, - }, - indent=2, - sort_keys=True, - ) - ) - return 0 - except ( - StarVLAError, - OSError, - ValueError, - RuntimeError, - KeyError, - json.JSONDecodeError, - ) as exc: - print(f"error: {exc}", file=sys.stderr) - return 1 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/tools/hf2gguf/starvla/generate_starvla_qwen25_groot_golden.py b/tools/hf2gguf/starvla/generate_starvla_qwen25_groot_golden.py deleted file mode 100644 index 0c6c461..0000000 --- a/tools/hf2gguf/starvla/generate_starvla_qwen25_groot_golden.py +++ /dev/null @@ -1,1142 +0,0 @@ -#!/usr/bin/env python3 -"""Generate a fixed-noise local-Python oracle for official Qwen2.5-VL GR00T. - -The oracle is intentionally separate from the Qwen3-VL GR00T schema. The -released Qwen2.5 model conditions its flow head on hidden tuple entry 36, -which is the final ``result_norm`` tensor, while the Qwen3 model uses the raw -outer ``l_out-35`` recorder boundary. -""" - -from __future__ import annotations - -import argparse -import datetime as dt -import gc -import json -import os -import random -import shutil -import sys -import tempfile -from pathlib import Path -from typing import Any, Iterable, Mapping, Sequence - -import numpy as np - - -TOOLS_DIR = Path(__file__).resolve().parent -if str(TOOLS_DIR) not in sys.path: - sys.path.insert(0, str(TOOLS_DIR)) - -from generate_starvla_oft_golden import ( # noqa: E402 - _array_record, - _assert_module_origin, - _canonical_json, - _configure_determinism, - _distribution_version, - _ensure_regular_file, - _image_pixel_sha256, - _require_isolated_python, - _runtime_record, - _sha256_bytes, - _tensor_to_array, - validate_runtime_versions, -) -from generate_starvla_qwen25_oft_golden import ( # noqa: E402 - EXPECTED_QWEN_VL_UTILS_VERSION, - _config_only_qwen25_bootstrap, - _official_qwen25_alias, - _qwen_asset_records, - _verify_clean_source, -) -from starvla_checkpoint import ( # noqa: E402 - DEFAULT_CATALOG, - StarVLAError, - get_variant, - load_catalog, - official_bundle_uuid, - sha256_file, - verify_catalog_files, - verify_checkpoint_file, -) - - -SCHEMA_VERSION = 1 -GOLDEN_KIND = "starvla_qwen25_groot_local_pt_python_oracle" -MODEL_TYPE = "starvla" -VARIANT = "qwen25_groot" -BACKBONE = "qwen2_5_vl" -ACTION_RELATIVE_L2_LIMIT = 0.03 -SEED = 0 - -OFFICIAL_CHECKPOINT_REPO_ID = "StarVLA/Qwen-GR00T-Bridge-RT-1" -OFFICIAL_CHECKPOINT_REVISION = "5ebc661ba38b29c28f20fff6574801e6f49f3466" -OFFICIAL_CHECKPOINT_FILENAME = "steps_30000_pytorch_model.pt" -OFFICIAL_CHECKPOINT_SIZE = 8_456_891_339 -OFFICIAL_CHECKPOINT_SHA256 = "9646da2ae0b32589a75c8cc88fae96c93c5d269b69fd7a29200744936e01d96f" -OFFICIAL_QWEN_REPO_ID = "StarVLA/Qwen2.5-VL-3B-Instruct-Action" -OFFICIAL_QWEN_REVISION = "ce86bd9a53416527b8361e8dfc47316288ffa110" -OFFICIAL_STARVLA_REPO_ID = "starVLA/starVLA" -OFFICIAL_STARVLA_REVISION = "631aae02afe6d95876e923ff518e8ff2ab9a2f88" - -EXPECTED_ACTION_HORIZON = 16 -EXPECTED_ACTION_DIM = 7 -EXPECTED_QWEN_HIDDEN_DIM = 2048 -EXPECTED_QWEN_LAYER_COUNT = 36 -EXPECTED_HIDDEN_TUPLE_INDEX = 36 -EXPECTED_DIT_WIDTH = 768 -EXPECTED_DIT_OUTPUT_DIM = 1024 -EXPECTED_DIT_BLOCK_COUNT = 16 -EXPECTED_FUTURE_TOKEN_COUNT = 32 -EXPECTED_TIMESTEP_IDS = [0, 250, 500, 750] -EXPECTED_COT_TEMPLATE = ( - "Your task is {instruction}. To identify the key objects for your task. " - "Locate their bounding boxes in [x1,y1,x2,y2] format." -) -ACTION_TOKEN_ID_MIN = 151665 -ACTION_TOKEN_ID_MAX = 153712 -ACTION_TOKEN_COUNT = 2048 -UNNORM_KEYS = ("oxe_bridge", "oxe_rt1") - -PINNED_SOURCE_FILES = { - "starVLA/model/framework/VLM4A/QwenGR00T.py": - "645d99d8d6a8daaccb7bb6e3211971b5cc7396d39968b0e9c20c3894d6883249", - "starVLA/model/modules/action_model/GR00T_ActionHeader.py": - "a01c7ca048589835a23bf46cf670275dfa643a1fb2da0bafd14654e1a57236e5", - "starVLA/model/modules/action_model/flow_matching_head/cross_attention_dit.py": - "c18d2e128dddcd67dc88c4fb178c99d7ceb7ea40d40ea9622b120151b81db359", - "starVLA/model/modules/vlm/QWen2_5.py": - "296a6b22859517ed9c302bc7c4e1c3362690e1da12ec8ad9a26b8a90d25dabec", - "deployment/model_server/policy_norm_processor.py": - "3fd280c8f5072943fad6809dd5705cb713007c10d7240d2a23e3dadcd3963d2a", -} - - -def _load_json_object(path: Path, *, label: str) -> dict[str, Any]: - _ensure_regular_file(path, label=label) - try: - value = json.loads(path.read_text(encoding="utf-8")) - except (OSError, UnicodeError, json.JSONDecodeError) as exc: - raise StarVLAError(f"failed to parse {label} {path}: {exc}") from exc - if not isinstance(value, dict): - raise StarVLAError(f"{label} root must be an object") - return value - - -def verify_source_semantics(source_dir: Path) -> dict[str, Any]: - actual: dict[str, str] = {} - for relative, expected in PINNED_SOURCE_FILES.items(): - path = source_dir / relative - _ensure_regular_file(path, label=f"pinned source {relative}") - digest = sha256_file(path) - if digest != expected: - raise StarVLAError( - f"pinned Qwen2.5 GR00T source SHA256 mismatch for {relative}: " - f"expected {expected}, got {digest}" - ) - actual[relative] = digest - - framework_source = ( - source_dir / "starVLA/model/framework/VLM4A/QwenGR00T.py" - ).read_text(encoding="utf-8") - action_source = ( - source_dir / "starVLA/model/modules/action_model/GR00T_ActionHeader.py" - ).read_text(encoding="utf-8") - required = ( - "last_hidden = qwenvl_outputs.hidden_states[-1]", - "backbone_attention_mask = backbone_attention_mask.to(dtype=torch.bool)", - "last_hidden, state, encoder_attention_mask=backbone_attention_mask", - "dtype=vl_embs.dtype", - "t_cont = t / float(num_steps)", - "t_discretized = int(t_cont * self.num_timestep_buckets)", - "actions = actions + dt * pred_velocity", - ) - combined = framework_source + "\n" + action_source - missing = [fragment for fragment in required if fragment not in combined] - if missing: - raise StarVLAError( - f"pinned Qwen2.5 GR00T source semantics probe failed: {missing!r}" - ) - return { - "files": actual, - "hidden_selection": "qwenvl_outputs.hidden_states[-1]", - "hidden_tuple_index": EXPECTED_HIDDEN_TUPLE_INDEX, - "hidden_tap": "result_norm", - "initial_noise": "torch.randn(dtype=vl_embs.dtype)", - "timestep_ids": EXPECTED_TIMESTEP_IDS, - "euler_update": "actions = actions + dt * pred_velocity", - } - - -def _validate_catalog_identity(catalog: Mapping[str, Any]) -> tuple[dict[str, Any], dict[str, Any]]: - variant = get_variant(catalog, VARIANT) - qwen_key = variant.get("qwen_asset") - qwen = catalog.get("shared_assets", {}).get(qwen_key) - if not isinstance(qwen, dict): - raise StarVLAError(f"catalog variant {VARIANT} has no Qwen action asset") - expected_variant = { - "repo_id": OFFICIAL_CHECKPOINT_REPO_ID, - "revision": OFFICIAL_CHECKPOINT_REVISION, - } - for key, expected in expected_variant.items(): - if variant.get(key) != expected: - raise StarVLAError( - f"catalog {VARIANT}.{key} must be {expected!r}, got {variant.get(key)!r}" - ) - expected_checkpoint = { - "path": f"checkpoints/{OFFICIAL_CHECKPOINT_FILENAME}", - "size": OFFICIAL_CHECKPOINT_SIZE, - "sha256": OFFICIAL_CHECKPOINT_SHA256, - } - if variant.get("checkpoint") != expected_checkpoint: - raise StarVLAError("catalog Qwen2.5 GR00T checkpoint identity drifted") - if ( - qwen.get("repo_id") != OFFICIAL_QWEN_REPO_ID - or qwen.get("revision") != OFFICIAL_QWEN_REVISION - ): - raise StarVLAError("catalog Qwen2.5 action-tokenizer identity drifted") - if catalog.get("source_revisions", {}).get("starvla") != OFFICIAL_STARVLA_REVISION: - raise StarVLAError("catalog StarVLA source revision drifted") - return variant, qwen - - -def validate_action_tokenizer_assets(qwen_dir: Path) -> dict[str, Any]: - config = _load_json_object(qwen_dir / "config.json", label="Qwen2.5 action config") - text_config = config.get("text_config") - actual = { - "model_type": config.get("model_type"), - "hidden_size": config.get("hidden_size"), - "text_hidden_size": text_config.get("hidden_size") if isinstance(text_config, dict) else None, - "layer_count": text_config.get("num_hidden_layers") if isinstance(text_config, dict) else None, - "vocab_size": text_config.get("vocab_size") if isinstance(text_config, dict) else None, - } - expected = { - "model_type": BACKBONE, - "hidden_size": EXPECTED_QWEN_HIDDEN_DIM, - "text_hidden_size": EXPECTED_QWEN_HIDDEN_DIM, - "layer_count": EXPECTED_QWEN_LAYER_COUNT, - "vocab_size": ACTION_TOKEN_ID_MAX + 1, - } - if actual != expected: - raise StarVLAError(f"unexpected Qwen2.5 action model config: {actual}") - - token_map = _load_json_object( - qwen_dir / "added_token_id_map.json", - label="Qwen2.5 action token map", - ) - expected_map = { - f"": ACTION_TOKEN_ID_MIN + index - for index in range(ACTION_TOKEN_COUNT) - } - if token_map != expected_map: - raise StarVLAError( - "Qwen2.5 action tokenizer must contain the contiguous " - "2048-token range 151665..153712" - ) - assets, assets_sha256 = _qwen_asset_records(qwen_dir) - return { - "repo_id": OFFICIAL_QWEN_REPO_ID, - "revision": OFFICIAL_QWEN_REVISION, - "action_token_count": ACTION_TOKEN_COUNT, - "action_token_id_min": ACTION_TOKEN_ID_MIN, - "action_token_id_max": ACTION_TOKEN_ID_MAX, - "assets": assets, - "assets_sha256": assets_sha256, - } - - -def _validate_effective_config(config: Mapping[str, Any]) -> None: - try: - framework = config["framework"] - action = framework["action_model"] - diffusion = action["diffusion_model_cfg"] - vla = config["datasets"]["vla_data"] - except (KeyError, TypeError) as exc: - raise StarVLAError("effective Qwen2.5 GR00T config is incomplete") from exc - actual = { - "framework_py": framework.get("framework_py"), - "action_model_type": action.get("action_model_type"), - "action_horizon": action.get("action_horizon"), - "action_dim": action.get("action_dim"), - "state_dim": action.get("state_dim"), - "steps": action.get("num_inference_timesteps"), - "buckets": action.get("num_timestep_buckets"), - "future_tokens": action.get("num_target_vision_tokens"), - "layers": diffusion.get("num_layers"), - "cross_dim": diffusion.get("cross_attention_dim"), - "output_dim": diffusion.get("output_dim"), - "interleave": diffusion.get("interleave_self_attention"), - "obs_image_size": vla.get("obs_image_size"), - "obs": vla.get("obs"), - "data_mix": vla.get("data_mix"), - "cot": vla.get("CoT_prompt"), - } - expected = { - "framework_py": "QwenFM", - "action_model_type": "DiT-B", - "action_horizon": EXPECTED_ACTION_HORIZON, - "action_dim": EXPECTED_ACTION_DIM, - "state_dim": EXPECTED_ACTION_DIM, - "steps": 4, - "buckets": 1000, - "future_tokens": EXPECTED_FUTURE_TOKEN_COUNT, - "layers": EXPECTED_DIT_BLOCK_COUNT, - "cross_dim": EXPECTED_QWEN_HIDDEN_DIM, - "output_dim": EXPECTED_DIT_OUTPUT_DIM, - "interleave": True, - "obs_image_size": None, - "obs": ["image_0"], - "data_mix": "bridge_rt_1", - "cot": EXPECTED_COT_TEMPLATE, - } - if actual != expected: - raise StarVLAError(f"unexpected effective Qwen2.5 GR00T config: {actual}") - - -def groot_normalization_contract( - norm_stats: Mapping[str, Any], unnorm_key: str -) -> dict[str, Any]: - if unnorm_key not in UNNORM_KEYS: - raise StarVLAError( - f"Qwen2.5 GR00T unnorm_key must be one of {list(UNNORM_KEYS)}, " - f"got {unnorm_key!r}" - ) - profile = norm_stats.get(unnorm_key) - action = profile.get("action") if isinstance(profile, Mapping) else None - if not isinstance(action, Mapping): - raise StarVLAError(f"dataset statistics has no {unnorm_key}.action object") - try: - q01 = np.asarray(action["q01"], dtype=np.float32) - q99 = np.asarray(action["q99"], dtype=np.float32) - mask = np.asarray(action["mask"], dtype=np.bool_) - except (KeyError, TypeError, ValueError) as exc: - raise StarVLAError(f"invalid GR00T action statistics for {unnorm_key}: {exc}") from exc - if q01.shape != (7,) or q99.shape != (7,) or mask.shape != (7,): - raise StarVLAError("Qwen2.5 GR00T action statistics must be 7D") - if not np.isfinite(q01).all() or not np.isfinite(q99).all(): - raise StarVLAError("Qwen2.5 GR00T action statistics must be finite") - if np.any(q99[mask] <= q01[mask]): - raise StarVLAError("Qwen2.5 GR00T masked q99 values must exceed q01") - return { - "stats_key": unnorm_key, - "runtime_robot_profile": unnorm_key, - "implementation": "official_PolicyNormProcessor_ComposedModalityTransform", - "q01": q01.tolist(), - "q99": q99.tolist(), - "mask": mask.tolist(), - } - - -def validate_local_inputs( - *, - checkpoint_root: Path, - checkpoint: Path | None, - qwen_model: Path | None, - source_dir: Path, - catalog_path: Path = DEFAULT_CATALOG, -) -> dict[str, Any]: - catalog = load_catalog(catalog_path) - variant, qwen = _validate_catalog_identity(catalog) - checkpoint_root = checkpoint_root.resolve() - policy_dir = checkpoint_root / "sources" / variant["directory"] / variant["revision"] - qwen_dir = checkpoint_root / "sources" / qwen["directory"] / qwen["revision"] - checkpoint_path = policy_dir / variant["checkpoint"]["path"] - if checkpoint is not None and checkpoint.resolve() != checkpoint_path.resolve(): - raise StarVLAError( - f"Qwen2.5 GR00T checkpoint must be the catalog path {checkpoint_path}" - ) - if qwen_model is not None and qwen_model.resolve() != qwen_dir.resolve(): - raise StarVLAError( - f"Qwen2.5 action processor must be the catalog path {qwen_dir}" - ) - - verify_catalog_files(policy_dir, variant) - verify_catalog_files(qwen_dir, qwen) - tokenizer = validate_action_tokenizer_assets(qwen_dir) - source_dir = source_dir.resolve() - revision = _verify_clean_source(source_dir, OFFICIAL_STARVLA_REVISION) - source_probe = verify_source_semantics(source_dir) - - sidecar = Path(f"{checkpoint_path}.aria2") - checkpoint_ready = ( - checkpoint_path.is_file() - and not checkpoint_path.is_symlink() - and not sidecar.exists() - ) - if checkpoint_ready: - verify_checkpoint_file(checkpoint_path, variant) - - config_yaml = policy_dir / "config.yaml" - dataset_statistics = policy_dir / "dataset_statistics.json" - norm_stats = _load_json_object( - dataset_statistics, label="Qwen2.5 GR00T dataset statistics" - ) - if set(norm_stats) != set(UNNORM_KEYS): - raise StarVLAError( - f"unexpected Qwen2.5 GR00T normalization profiles: {sorted(norm_stats)}" - ) - for key in UNNORM_KEYS: - groot_normalization_contract(norm_stats, key) - - try: - import yaml - - config = yaml.safe_load(config_yaml.read_text(encoding="utf-8")) - except (ImportError, OSError, UnicodeError, ValueError) as exc: - raise StarVLAError(f"failed to load Qwen2.5 GR00T config.yaml: {exc}") from exc - if not isinstance(config, dict): - raise StarVLAError("Qwen2.5 GR00T config.yaml root must be an object") - _validate_effective_config(config) - return { - "catalog": catalog, - "catalog_path": catalog_path.resolve(), - "variant": variant, - "qwen": qwen, - "policy_dir": policy_dir.resolve(), - "qwen_dir": qwen_dir.resolve(), - "checkpoint": checkpoint_path.resolve(), - "checkpoint_ready": checkpoint_ready, - "config_yaml": config_yaml.resolve(), - "dataset_statistics": dataset_statistics.resolve(), - "norm_stats": norm_stats, - "config": config, - "source_dir": source_dir, - "source_revision": revision, - "source_probe": source_probe, - "tokenizer": tokenizer, - } - - -def validate_processor_contract(qwen_dir: Path) -> dict[str, Any]: - import transformers - - processor = transformers.AutoProcessor.from_pretrained( - qwen_dir, - local_files_only=True, - trust_remote_code=False, - ) - tokenizer = processor.tokenizer - first = tokenizer("", add_special_tokens=False)["input_ids"] - last = tokenizer("", add_special_tokens=False)["input_ids"] - actual = { - "processor_class": type(processor).__name__, - "image_processor_class": type(processor.image_processor).__name__, - "tokenizer_length": len(tokenizer), - "first_action_token_ids": first, - "last_action_token_ids": last, - "padding_side": tokenizer.padding_side, - } - expected = { - "processor_class": "Qwen2_5_VLProcessor", - "image_processor_class": "Qwen2VLImageProcessorFast", - "tokenizer_length": ACTION_TOKEN_ID_MAX + 1, - "first_action_token_ids": [ACTION_TOKEN_ID_MIN], - "last_action_token_ids": [ACTION_TOKEN_ID_MAX], - "padding_side": "right", - } - # The official wrapper switches padding to left immediately after loading. - if actual != expected: - raise StarVLAError(f"unexpected Qwen2.5 action processor contract: {actual}") - return {**actual, "wrapper_padding_side": "left"} - - -def load_official_framework(paths: Mapping[str, Any], *, device: str) -> tuple[Any, dict[str, Any]]: - import torch - import transformers - - if not paths["checkpoint_ready"]: - raise StarVLAError( - f"official Qwen2.5 GR00T checkpoint is absent or incomplete: {paths['checkpoint']}" - ) - source_dir = Path(paths["source_dir"]) - if any(name == "starVLA" or name.startswith("starVLA.") for name in sys.modules): - raise StarVLAError("starVLA was imported before pinned-source verification") - sys.path.insert(0, str(source_dir)) - try: - from starVLA.model.framework import base_framework, share_tools - from starVLA.model.framework.VLM4A import QwenGR00T - - for module in (base_framework, share_tools, QwenGR00T): - _assert_module_origin(module, source_dir) - config, norm_stats = share_tools.read_mode_config(str(paths["checkpoint"])) - _validate_effective_config(config) - with _official_qwen25_alias(Path(paths["qwen_dir"])) as qwen_alias: - config = base_framework.merge_config_overrides( - config, - [ - f"framework.qwenvl.base_vlm={qwen_alias}", - "framework.qwenvl.attn_implementation=sdpa", - ], - ) - cfg = share_tools.dict_to_namespace(config) - cfg.trainer.pretrained_checkpoint = None - with _config_only_qwen25_bootstrap( - torch, transformers, Path(paths["qwen_dir"]) - ): - framework = QwenGR00T.Qwen_GR00T(cfg) - - try: - state = torch.load( - paths["checkpoint"], map_location="cpu", mmap=True, weights_only=True - ) - except TypeError: - state = torch.load(paths["checkpoint"], map_location="cpu", weights_only=True) - if not isinstance(state, Mapping) or not state: - raise StarVLAError("official Qwen2.5 GR00T checkpoint has no state_dict") - framework.load_state_dict(state, strict=True) - del state - gc.collect() - framework.norm_stats = norm_stats - - action_model = framework.action_model - if type(framework).__name__ != "Qwen_GR00T": - raise StarVLAError(f"unexpected official framework class: {type(framework).__name__}") - if int(framework.action_horizon) != EXPECTED_ACTION_HORIZON: - raise StarVLAError("official Qwen2.5 GR00T action horizon changed") - if len(action_model.model.transformer_blocks) != EXPECTED_DIT_BLOCK_COUNT: - raise StarVLAError("official Qwen2.5 GR00T DiT block count changed") - hidden_size = int(framework.qwen_vl_interface.model.config.hidden_size) - if hidden_size != EXPECTED_QWEN_HIDDEN_DIM: - raise StarVLAError(f"unexpected Qwen2.5 hidden size: {hidden_size}") - tokenizer = framework.qwen_vl_interface.processor.tokenizer - if ( - len(tokenizer) != ACTION_TOKEN_ID_MAX + 1 - or tokenizer.convert_tokens_to_ids("") != ACTION_TOKEN_ID_MIN - or tokenizer.convert_tokens_to_ids("") != ACTION_TOKEN_ID_MAX - ): - raise StarVLAError("official framework did not load the pinned action tokenizer") - - qwen_dtypes = {parameter.dtype for parameter in framework.qwen_vl_interface.parameters()} - policy_dtypes = {parameter.dtype for parameter in action_model.parameters()} - if qwen_dtypes != {torch.bfloat16} or policy_dtypes != {torch.float32}: - raise StarVLAError( - "official Qwen2.5 GR00T dtype boundary changed: " - f"qwen={qwen_dtypes}, policy={policy_dtypes}" - ) - return framework.to(device).eval(), config - finally: - if sys.path and sys.path[0] == str(source_dir): - del sys.path[0] - - -def _first_tensor(value: Any) -> Any: - if isinstance(value, (tuple, list)): - if not value: - raise StarVLAError("decoder layer returned an empty tuple") - return value[0] - return value - - -def run_official_forward( - framework: Any, - *, - images: Sequence[Any], - task: str, - seed: int = SEED, -) -> dict[str, Any]: - """Run the official framework and capture its result_norm/fixed-noise boundary.""" - - import torch - - captures: dict[str, Any] = {} - qwen = framework.qwen_vl_interface - action_model = framework.action_model - language_model = qwen.model.model.language_model - handles = [] - original_build = qwen.build_qwenvl_inputs - original_policy = action_model.predict_action - - def capture_build(*args: Any, **kwargs: Any): - batch_images = kwargs.get("images", args[0] if args else None) - instructions = kwargs.get("instructions", args[1] if len(args) > 1 else None) - captures["processed_images"] = list(batch_images[0]) - captures["framework_instructions"] = list(instructions) - output = original_build(*args, **kwargs) - captures["qwen_inputs"] = { - key: value.detach() - for key, value in output.items() - if isinstance(value, torch.Tensor) - } - return output - - def capture_outer(_module: Any, _inputs: Any, output: Any): - hidden = getattr(output, "hidden_states", None) - if hidden is None or len(hidden) != EXPECTED_QWEN_LAYER_COUNT + 1: - raise StarVLAError( - "official Qwen2.5 outer output did not expose 37 hidden tuple entries" - ) - captures["outer_final"] = hidden[EXPECTED_HIDDEN_TUPLE_INDEX].detach().clone() - - def capture_policy(*args: Any, **kwargs: Any): - conditioning = args[0] if args else kwargs.get("vl_embs") - state = args[1] if len(args) > 1 else kwargs.get("state") - mask = kwargs.get("encoder_attention_mask", args[2] if len(args) > 2 else None) - if state is not None: - raise StarVLAError("official Qwen2.5 GR00T unexpectedly used state") - captures["policy_conditioning"] = conditioning.detach().clone() - captures["policy_attention_mask"] = mask.detach().clone() - original_randn = torch.randn - - def capture_randn(*randn_args: Any, **randn_kwargs: Any): - value = original_randn(*randn_args, **randn_kwargs) - if "initial_noise" in captures: - raise StarVLAError("official Qwen2.5 GR00T sampled noise more than once") - captures["initial_noise"] = value.detach().clone() - return value - - torch.randn = capture_randn - try: - output = original_policy(*args, **kwargs) - finally: - torch.randn = original_randn - captures["raw_policy"] = output.detach().clone() - return output - - handles.append( - language_model.layers[-1].register_forward_hook( - lambda _m, _i, output: captures.__setitem__( - "raw_l_out_35", _first_tensor(output).detach().clone() - ) - ) - ) - handles.append( - language_model.norm.register_forward_hook( - lambda _m, _i, output: captures.__setitem__( - "result_norm", output.detach().clone() - ) - ) - ) - handles.append(qwen.model.register_forward_hook(capture_outer)) - qwen.build_qwenvl_inputs = capture_build - action_model.predict_action = capture_policy - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - try: - result = framework.predict_action( - examples=[{"image": list(images), "lang": task}] - ) - finally: - for handle in handles: - handle.remove() - qwen.build_qwenvl_inputs = original_build - action_model.predict_action = original_policy - - required = { - "processed_images", - "framework_instructions", - "qwen_inputs", - "raw_l_out_35", - "result_norm", - "outer_final", - "policy_conditioning", - "policy_attention_mask", - "initial_noise", - "raw_policy", - } - missing = sorted(required - set(captures)) - if missing: - raise StarVLAError(f"Qwen2.5 GR00T instrumentation missed: {missing}") - if captures["framework_instructions"] != [task]: - raise StarVLAError("official framework changed the input instruction") - if len(captures["processed_images"]) != len(images) or any( - actual.mode != expected.mode - or actual.size != expected.size - or actual.tobytes() != expected.tobytes() - for actual, expected in zip(captures["processed_images"], images, strict=True) - ): - raise StarVLAError("official Qwen2.5 GR00T unexpectedly pre-resized the image") - if not torch.equal(captures["outer_final"], captures["result_norm"]): - raise StarVLAError("Qwen2.5 hidden tuple entry 36 is not result_norm") - if torch.equal(captures["outer_final"], captures["raw_l_out_35"]): - raise StarVLAError("Qwen2.5 result_norm unexpectedly equals raw l_out-35") - if not torch.equal(captures["policy_conditioning"], captures["result_norm"]): - raise StarVLAError("Qwen2.5 GR00T policy did not receive result_norm") - mask = captures["qwen_inputs"].get("attention_mask") - if mask is None or not torch.equal( - captures["policy_attention_mask"], mask.to(dtype=torch.bool) - ): - raise StarVLAError("Qwen2.5 GR00T policy mask is not the full boolean Qwen mask") - if captures["policy_attention_mask"].dtype != torch.bool: - raise StarVLAError("Qwen2.5 GR00T policy attention mask must be bool") - if captures["result_norm"].dtype != torch.bfloat16: - raise StarVLAError("Qwen2.5 result_norm boundary must be BF16") - if captures["initial_noise"].dtype != torch.bfloat16: - raise StarVLAError("Qwen2.5 GR00T initial noise must be sampled as BF16") - if tuple(captures["initial_noise"].shape) != (1, 16, 7): - raise StarVLAError("Qwen2.5 GR00T initial noise shape must be [1,16,7]") - - normalized = np.asarray(result.get("normalized_actions"), dtype=np.float32) - raw_policy, _ = _tensor_to_array(captures["raw_policy"]) - if normalized.shape != (1, EXPECTED_ACTION_HORIZON, EXPECTED_ACTION_DIM): - raise StarVLAError(f"unexpected Qwen2.5 GR00T output shape: {normalized.shape}") - if normalized.shape != raw_policy.shape or not np.array_equal(normalized, raw_policy): - raise StarVLAError("normalized actions differ from captured GR00T policy output") - if not np.isfinite(normalized).all(): - raise StarVLAError("official Qwen2.5 GR00T produced non-finite actions") - captures["normalized_actions"] = np.ascontiguousarray(normalized) - return captures - - -def _load_images(image_paths: Iterable[Path]) -> tuple[list[Any], list[dict[str, Any]]]: - from PIL import Image - - images: list[Any] = [] - records: list[dict[str, Any]] = [] - for path in image_paths: - path = path.resolve() - _ensure_regular_file(path, label="Qwen2.5 GR00T input image") - try: - with Image.open(path) as opened: - opened.load() - image = opened.convert("RGB") - except (OSError, ValueError) as exc: - raise StarVLAError(f"failed to decode input image {path}: {exc}") from exc - images.append(image) - records.append( - { - "source_path": str(path), - "source_size": path.stat().st_size, - "source_sha256": sha256_file(path), - "decoded_mode": image.mode, - "decoded_size": list(image.size), - "decoded_pixel_sha256": _image_pixel_sha256(image), - } - ) - if len(images) != 1: - raise StarVLAError("official Qwen2.5 GR00T oracle requires exactly one image") - return images, records - - -def _render_model_prompt(framework: Any, images: Sequence[Any], task: str) -> str: - model_instruction = EXPECTED_COT_TEMPLATE.replace("{instruction}", task) - messages = [{ - "role": "user", - "content": [ - *({"type": "image", "image": image} for image in images), - {"type": "text", "text": model_instruction}, - ], - }] - rendered = framework.qwen_vl_interface.processor.apply_chat_template( - messages, tokenize=False, add_generation_prompt=True - ) - if not isinstance(rendered, str): - raise StarVLAError("Qwen2.5 processor returned a non-string prompt") - return rendered - - -def _build_arrays( - captures: Mapping[str, Any], unnormalized: np.ndarray -) -> tuple[dict[str, np.ndarray], dict[str, Any]]: - required_inputs = ("input_ids", "attention_mask", "image_grid_thw") - missing = [name for name in required_inputs if name not in captures["qwen_inputs"]] - if missing: - raise StarVLAError(f"Qwen2.5 processor outputs are missing: {missing}") - values = { - "input_ids": captures["qwen_inputs"]["input_ids"], - "attention_mask": captures["qwen_inputs"]["attention_mask"], - "image_grid_thw": captures["qwen_inputs"]["image_grid_thw"], - "raw_l_out_35_diagnostic": captures["raw_l_out_35"], - "result_norm": captures["result_norm"], - "initial_noise": captures["initial_noise"], - "normalized_actions": captures["normalized_actions"], - "unnormalized_actions": np.ascontiguousarray(unnormalized, dtype=np.float32), - } - arrays: dict[str, np.ndarray] = {} - records: dict[str, Any] = {} - for name, value in values.items(): - if isinstance(value, np.ndarray): - array = np.ascontiguousarray(value) - source_dtype = None - else: - array, source_dtype = _tensor_to_array(value) - arrays[name] = array - records[name] = _array_record(array, source_dtype=source_dtype) - - token_count = arrays["input_ids"].shape[1] - expected_shapes = { - "input_ids": (1, token_count), - "attention_mask": (1, token_count), - "image_grid_thw": (1, 3), - "raw_l_out_35_diagnostic": (1, token_count, EXPECTED_QWEN_HIDDEN_DIM), - "result_norm": (1, token_count, EXPECTED_QWEN_HIDDEN_DIM), - "initial_noise": (1, EXPECTED_ACTION_HORIZON, EXPECTED_ACTION_DIM), - "normalized_actions": (1, EXPECTED_ACTION_HORIZON, EXPECTED_ACTION_DIM), - "unnormalized_actions": (1, EXPECTED_ACTION_HORIZON, EXPECTED_ACTION_DIM), - } - for name, expected in expected_shapes.items(): - if arrays[name].shape != expected: - raise StarVLAError( - f"Qwen2.5 GR00T {name} shape must be {expected}, got {arrays[name].shape}" - ) - return arrays, records - - -def write_golden( - *, - output_dir: Path, - paths: Mapping[str, Any], - framework: Any, - image_paths: Sequence[Path], - source_image_records: Sequence[Mapping[str, Any]], - task: str, - unnorm_key: str, - captures: Mapping[str, Any], - unnormalized: np.ndarray, -) -> Path: - import torch - import transformers - - output_dir = output_dir.resolve() - if output_dir.exists(): - raise StarVLAError(f"golden output directory already exists: {output_dir}") - output_dir.parent.mkdir(parents=True, exist_ok=True) - arrays, records = _build_arrays(captures, unnormalized) - model_instruction = EXPECTED_COT_TEMPLATE.replace("{instruction}", task) - rendered_prompt = _render_model_prompt( - framework, captures["processed_images"], task - ) - identity = { - "schema_version": SCHEMA_VERSION, - "kind": GOLDEN_KIND, - "checkpoint_sha256": OFFICIAL_CHECKPOINT_SHA256, - "starvla_revision": paths["source_revision"], - "qwen_revision": OFFICIAL_QWEN_REVISION, - "task": task, - "unnorm_key": unnorm_key, - "images": [record["source_sha256"] for record in source_image_records], - "initial_noise_sha256": records["initial_noise"]["sha256"], - } - golden_id = _sha256_bytes(_canonical_json(identity)) - - with tempfile.TemporaryDirectory( - prefix=f".{output_dir.name}.", dir=output_dir.parent - ) as temporary: - staging = Path(temporary) - inputs_dir = staging / "inputs" - inputs_dir.mkdir() - image_records: list[dict[str, Any]] = [] - for index, (source_path, source_record) in enumerate( - zip(image_paths, source_image_records, strict=True) - ): - suffix = source_path.suffix.lower() or ".img" - artifact = inputs_dir / f"image-{index:02d}{suffix}" - shutil.copyfile(source_path, artifact) - image_records.append({ - **source_record, - "artifact": artifact.relative_to(staging).as_posix(), - "artifact_size": artifact.stat().st_size, - "artifact_sha256": sha256_file(artifact), - }) - - tensors_path = staging / "tensors.npz" - np.savez(tensors_path, **arrays) - noise_bytes = np.ascontiguousarray( - arrays["initial_noise"], dtype=" dict[str, Any]: - return { - "schema_version": SCHEMA_VERSION, - "kind": "starvla_qwen25_groot_preflight", - "variant": VARIANT, - "model_type": MODEL_TYPE, - "backbone": BACKBONE, - "checkpoint": str(paths["checkpoint"]), - "checkpoint_ready": paths["checkpoint_ready"], - "expected_checkpoint": { - "bundle_uuid": official_bundle_uuid( - paths["variant"], paths["catalog"] - ), - "repo_id": OFFICIAL_CHECKPOINT_REPO_ID, - "revision": OFFICIAL_CHECKPOINT_REVISION, - "filename": OFFICIAL_CHECKPOINT_FILENAME, - "size": OFFICIAL_CHECKPOINT_SIZE, - "sha256": OFFICIAL_CHECKPOINT_SHA256, - }, - "qwen": { - **paths["tokenizer"], - "processor": dict(processor), - }, - "conditioning": { - "hidden_tuple_index": EXPECTED_HIDDEN_TUPLE_INDEX, - "hidden_tap_name": "result_norm", - "hidden_size": EXPECTED_QWEN_HIDDEN_DIM, - }, - "action": { - "shape": [1, EXPECTED_ACTION_HORIZON, EXPECTED_ACTION_DIM], - "initial_noise_dtype": "bfloat16", - "timestep_ids": EXPECTED_TIMESTEP_IDS, - }, - "action_gate": { - "reference": "local_official_python_pt", - "metric": "full_tensor_global_relative_l2", - "operator": "<=", - "limit": ACTION_RELATIVE_L2_LIMIT, - "required_outputs": ["normalized_actions", "unnormalized_actions"], - }, - "source_probe": paths["source_probe"], - "effective_config_valid": True, - "golden_created": False, - } - - -def build_parser() -> argparse.ArgumentParser: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--checkpoint-root", type=Path, default=Path("ckpts/starvla")) - parser.add_argument("--checkpoint", type=Path) - parser.add_argument("--qwen-model", type=Path) - parser.add_argument( - "--starvla-source", - type=Path, - default=Path("ckpts/starvla/source/starvla"), - ) - parser.add_argument("--catalog", type=Path, default=DEFAULT_CATALOG) - parser.add_argument("--image", action="append", default=[], type=Path) - parser.add_argument("--task", default="grab the block.") - parser.add_argument("--unnorm-key", choices=UNNORM_KEYS, default="oxe_bridge") - parser.add_argument("--device", default="cuda:0") - parser.add_argument( - "--output-dir", - type=Path, - default=Path("goldens/starvla/qwen25-groot/bridge-grab-block"), - ) - parser.add_argument("--preflight", "--preflight-only", action="store_true") - return parser - - -def main(argv: Sequence[str] | None = None) -> int: - args = build_parser().parse_args(argv) - try: - _require_isolated_python() - import torch - import transformers - - validate_runtime_versions( - torch_version=torch.__version__, - torchvision_version=_distribution_version("torchvision"), - transformers_version=transformers.__version__, - numpy_version=np.__version__, - ) - if _distribution_version("qwen-vl-utils") != EXPECTED_QWEN_VL_UTILS_VERSION: - raise StarVLAError( - "qwen-vl-utils must be " - f"{EXPECTED_QWEN_VL_UTILS_VERSION} for the official oracle" - ) - _configure_determinism(torch, seed=SEED, device=args.device) - paths = validate_local_inputs( - checkpoint_root=args.checkpoint_root, - checkpoint=args.checkpoint, - qwen_model=args.qwen_model, - source_dir=args.starvla_source, - catalog_path=args.catalog, - ) - processor = validate_processor_contract(Path(paths["qwen_dir"])) - if args.preflight: - print(json.dumps( - _preflight_record(paths, processor), - allow_nan=False, - indent=2, - sort_keys=True, - )) - return 0 - if not paths["checkpoint_ready"]: - raise StarVLAError( - f"official Qwen2.5 GR00T checkpoint is not ready: {paths['checkpoint']}" - ) - if len(args.image) != 1: - raise StarVLAError("exactly one --image is required") - images, image_records = _load_images(args.image) - framework, _config = load_official_framework(paths, device=args.device) - captures = run_official_forward( - framework, images=images, task=args.task, seed=SEED - ) - - source_dir = Path(paths["source_dir"]) - sys.path.insert(0, str(source_dir)) - try: - from deployment.model_server import policy_norm_processor - - _assert_module_origin(policy_norm_processor, source_dir) - normalizer = policy_norm_processor.PolicyNormProcessor( - str(paths["checkpoint"]), unnorm_key=args.unnorm_key - ) - unnormalized = np.stack([ - normalizer.unapply_actions(captures["normalized_actions"][0]) - ]).astype(np.float32, copy=False) - finally: - if sys.path and sys.path[0] == str(source_dir): - del sys.path[0] - if unnormalized.shape != (1, EXPECTED_ACTION_HORIZON, EXPECTED_ACTION_DIM): - raise StarVLAError( - f"official Qwen2.5 GR00T unnormalized shape changed: {unnormalized.shape}" - ) - if not np.isfinite(unnormalized).all(): - raise StarVLAError("official Qwen2.5 GR00T unnormalized actions are non-finite") - - manifest = write_golden( - output_dir=args.output_dir, - paths=paths, - framework=framework, - image_paths=args.image, - source_image_records=image_records, - task=args.task, - unnorm_key=args.unnorm_key, - captures=captures, - unnormalized=unnormalized, - ) - print(manifest) - return 0 - except (StarVLAError, OSError, ValueError, KeyError, RuntimeError) as exc: - print(f"error: {exc}", file=sys.stderr) - return 1 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/tools/hf2gguf/starvla/generate_starvla_qwen25_oft_golden.py b/tools/hf2gguf/starvla/generate_starvla_qwen25_oft_golden.py deleted file mode 100644 index 87dfdfa..0000000 --- a/tools/hf2gguf/starvla/generate_starvla_qwen25_oft_golden.py +++ /dev/null @@ -1,986 +0,0 @@ -#!/usr/bin/env python3 -"""Generate a local-Python oracle from the official Qwen2.5-VL OFT .pt file. - -Unlike the Qwen3-VL catalog-driven oracle, this entry point binds the golden -directly to an explicitly supplied local checkpoint and processor directory. -The full StarVLA checkpoint supplies all model parameters; the Qwen directory -is used only for the model topology and processor/tokenizer assets. -""" - -from __future__ import annotations - -import argparse -import contextlib -import datetime as dt -import gc -import json -import os -import shutil -import sys -import tempfile -from pathlib import Path -from typing import Any, Iterable, Mapping, Sequence - -import numpy as np - - -TOOLS_DIR = Path(__file__).resolve().parent -if str(TOOLS_DIR) not in sys.path: - sys.path.insert(0, str(TOOLS_DIR)) - -from generate_starvla_oft_golden import ( # noqa: E402 - _array_record, - _assert_module_origin, - _canonical_json, - _configure_determinism, - _distribution_version, - _ensure_regular_file, - _image_pixel_sha256, - _require_isolated_python, - _runtime_record, - _sha256_bytes, - _tensor_to_array, - expected_framework_instruction, - expected_model_instruction, - select_action_positions, - validate_runtime_versions, -) -from starvla_checkpoint import StarVLAError, sha256_file # noqa: E402 - - -SCHEMA_VERSION = 1 -GOLDEN_KIND = "starvla_qwen25_oft_local_pt_python_oracle" -MODEL_TYPE = "starvla" -BACKBONE = "qwen2_5_vl" -ACTION_TOKEN = chr(0x1F50D) -ACTION_RELATIVE_L2_LIMIT = 0.03 -EXPECTED_QWEN_VL_UTILS_VERSION = "0.0.14" -LEGACY_UNNORM_PROFILES = { - "bridge_dataset": "oxe_bridge", - "fractal20220817_data": "oxe_rt1", -} - -OFFICIAL_CHECKPOINT_REPO_ID = "StarVLA/Qwen-OFT-Bridge-RT-1" -OFFICIAL_CHECKPOINT_REVISION = "11fa6440835ba3e912de43cfe8521043360ffc02" -OFFICIAL_CHECKPOINT_FILENAME = "steps_10000_pytorch_model.pt" -OFFICIAL_CHECKPOINT_SIZE = 8_215_912_766 -OFFICIAL_CHECKPOINT_SHA256 = "51fe8d22c8d57116c2f59c5fdb24323fa3411149e888b807edba99b8354e0861" -OFFICIAL_QWEN_REVISION = "66285546d2b821cf421d4f5eb2576359d3770cd3" -OFFICIAL_BUNDLE_UUID = "90d105ae-00fa-580c-8751-9f931e324c3b" - -QWEN_PROCESSOR_REQUIRED = { - "config.json", - "preprocessor_config.json", - "tokenizer_config.json", -} -QWEN_PROCESSOR_OPTIONAL = { - "added_tokens.json", - "chat_template.json", - "merges.txt", - "special_tokens_map.json", - "tokenizer.json", - "vocab.json", -} - - -def _run_git(source_dir: Path, *arguments: str) -> str: - import subprocess - - try: - result = subprocess.run( - ["git", "-C", str(source_dir), *arguments], - check=True, - capture_output=True, - text=True, - ) - except (OSError, subprocess.CalledProcessError) as exc: - raise StarVLAError(f"failed to inspect StarVLA checkout {source_dir}: {exc}") from exc - return result.stdout.strip() - - -def _verify_clean_source(source_dir: Path, expected_revision: str | None) -> str: - source_dir = source_dir.resolve() - if not (source_dir / ".git").is_dir(): - raise StarVLAError(f"StarVLA source is not a Git checkout: {source_dir}") - revision = _run_git(source_dir, "rev-parse", "HEAD") - if expected_revision is not None and revision != expected_revision: - raise StarVLAError( - f"StarVLA source revision mismatch: expected {expected_revision}, got {revision}" - ) - changes = _run_git(source_dir, "status", "--porcelain=v1", "--untracked-files=all") - if changes: - raise StarVLAError(f"StarVLA source checkout is not clean:\n{changes}") - return revision - - -def _load_json_object(path: Path, *, label: str) -> dict[str, Any]: - _ensure_regular_file(path, label=label) - try: - value = json.loads(path.read_text(encoding="utf-8")) - except (OSError, UnicodeError, json.JSONDecodeError) as exc: - raise StarVLAError(f"failed to parse {label} {path}: {exc}") from exc - if not isinstance(value, dict): - raise StarVLAError(f"{label} root must be an object") - return value - - -def _qwen_asset_records(qwen_dir: Path) -> tuple[list[dict[str, Any]], str]: - qwen_dir = qwen_dir.resolve() - missing = sorted(name for name in QWEN_PROCESSOR_REQUIRED if not (qwen_dir / name).is_file()) - if missing: - raise StarVLAError("Qwen2.5 processor directory is missing: " + ", ".join(missing)) - - tokenizer_files = {"tokenizer.json", "vocab.json"} - if not any((qwen_dir / name).is_file() for name in tokenizer_files): - raise StarVLAError("Qwen2.5 processor directory needs tokenizer.json or vocab.json") - - records: list[dict[str, Any]] = [] - for name in sorted(QWEN_PROCESSOR_REQUIRED | QWEN_PROCESSOR_OPTIONAL): - path = qwen_dir / name - if not path.exists(): - continue - _ensure_regular_file(path, label=f"Qwen2.5 asset {name}") - records.append( - { - "path": name, - "size": path.stat().st_size, - "sha256": sha256_file(path), - } - ) - digest = _sha256_bytes(_canonical_json(records)) - return records, digest - - -def validate_local_inputs( - *, - checkpoint: Path, - qwen_model: Path, - source_dir: Path, - expected_checkpoint_sha256: str = OFFICIAL_CHECKPOINT_SHA256, - expected_checkpoint_size: int = OFFICIAL_CHECKPOINT_SIZE, - expected_source_revision: str | None = None, -) -> dict[str, Any]: - checkpoint = checkpoint.resolve() - _ensure_regular_file(checkpoint, label="Qwen2.5 OFT checkpoint") - if checkpoint.suffix != ".pt": - raise StarVLAError(f"Qwen2.5 OFT reference checkpoint must be a .pt file: {checkpoint}") - if Path(f"{checkpoint}.aria2").exists(): - raise StarVLAError(f"checkpoint download is incomplete: {checkpoint}.aria2 exists") - actual_size = checkpoint.stat().st_size - if actual_size != expected_checkpoint_size: - raise StarVLAError( - f"checkpoint size mismatch: expected {expected_checkpoint_size}, got {actual_size}" - ) - actual_sha256 = sha256_file(checkpoint) - if actual_sha256 != expected_checkpoint_sha256: - raise StarVLAError( - f"checkpoint SHA256 mismatch: expected {expected_checkpoint_sha256}, got {actual_sha256}" - ) - if checkpoint.parent.name != "checkpoints" or len(checkpoint.parents) < 2: - raise StarVLAError( - "checkpoint must use the StarVLA layout /checkpoints/.pt" - ) - - run_dir = checkpoint.parents[1] - config_yaml = run_dir / "config.yaml" - dataset_statistics = run_dir / "dataset_statistics.json" - _ensure_regular_file(config_yaml, label="checkpoint config.yaml") - stats = _load_json_object(dataset_statistics, label="checkpoint dataset_statistics.json") - if not stats: - raise StarVLAError("checkpoint dataset_statistics.json must not be empty") - - qwen_model = qwen_model.resolve() - if not qwen_model.is_dir(): - raise StarVLAError(f"Qwen2.5 processor/model directory does not exist: {qwen_model}") - qwen_config_path = qwen_model / "config.json" - qwen_config = _load_json_object(qwen_config_path, label="Qwen2.5 config.json") - if qwen_config.get("model_type") != BACKBONE: - raise StarVLAError( - f"Qwen config model_type must be {BACKBONE!r}, got {qwen_config.get('model_type')!r}" - ) - hidden_size = qwen_config.get("hidden_size") - if hidden_size != 2048: - raise StarVLAError(f"official Qwen2.5-VL 3B hidden_size must be 2048, got {hidden_size!r}") - qwen_assets, qwen_assets_sha256 = _qwen_asset_records(qwen_model) - - source_dir = source_dir.resolve() - source_revision = _verify_clean_source(source_dir, expected_source_revision) - return { - "checkpoint": checkpoint, - "checkpoint_size": actual_size, - "checkpoint_sha256": actual_sha256, - "run_dir": run_dir, - "config_yaml": config_yaml, - "dataset_statistics": dataset_statistics, - "norm_stats": stats, - "qwen_dir": qwen_model, - "qwen_config": qwen_config, - "qwen_assets": qwen_assets, - "qwen_assets_sha256": qwen_assets_sha256, - "source_dir": source_dir, - "source_revision": source_revision, - } - - -def legacy_normalization_contract( - norm_stats: Mapping[str, Any], unnorm_key: str -) -> dict[str, Any]: - if unnorm_key not in LEGACY_UNNORM_PROFILES: - raise StarVLAError( - f"Qwen2.5 OFT unnorm_key must be one of {sorted(LEGACY_UNNORM_PROFILES)}, " - f"got {unnorm_key!r}" - ) - value = norm_stats.get(unnorm_key) - if not isinstance(value, Mapping) or not isinstance(value.get("action"), Mapping): - raise StarVLAError(f"dataset statistics has no {unnorm_key}.action object") - action = value["action"] - try: - q01 = np.asarray(action["q01"], dtype=np.float32) - q99 = np.asarray(action["q99"], dtype=np.float32) - mask = np.asarray(action["mask"], dtype=np.bool_) - except (KeyError, TypeError, ValueError) as exc: - raise StarVLAError(f"invalid legacy action statistics for {unnorm_key}: {exc}") from exc - if q01.shape != (7,) or q99.shape != (7,) or mask.shape != (7,): - raise StarVLAError( - f"legacy action statistics must be 7D, got {q01.shape}/{q99.shape}/{mask.shape}" - ) - if not np.isfinite(q01).all() or not np.isfinite(q99).all(): - raise StarVLAError("legacy Qwen2.5 OFT q01/q99 statistics must be finite") - if np.any(q99[mask] <= q01[mask]): - raise StarVLAError("legacy Qwen2.5 OFT masked q99 values must exceed q01") - return { - "stats_key": unnorm_key, - "runtime_robot_profile": LEGACY_UNNORM_PROFILES[unnorm_key], - "method": "q01_q99_masked_with_binary_unmasked_dimensions", - "binary_threshold": 0.5, - "q01": q01.tolist(), - "q99": q99.tolist(), - "mask": mask.tolist(), - } - - -def unnormalize_legacy_actions( - normalized: np.ndarray, - norm_stats: Mapping[str, Any], - unnorm_key: str, -) -> np.ndarray: - """Mirror the released checkpoint's q99 + binary action transform.""" - - contract = legacy_normalization_contract(norm_stats, unnorm_key) - actions = np.ascontiguousarray(normalized, dtype=np.float32) - if actions.ndim != 3 or actions.shape[0] != 1 or actions.shape[-1] != 7: - raise StarVLAError( - f"normalized Qwen2.5 OFT actions must have shape [1,T,7], got {actions.shape}" - ) - q01 = np.asarray(contract["q01"], dtype=np.float32) - q99 = np.asarray(contract["q99"], dtype=np.float32) - mask = np.asarray(contract["mask"], dtype=np.bool_) - result = np.empty_like(actions) - result[..., mask] = ( - (actions[..., mask] + np.float32(1.0)) - / np.float32(2.0) - * (q99[mask] - q01[mask]) - + q01[mask] - ) - result[..., ~mask] = ( - actions[..., ~mask] > np.float32(contract["binary_threshold"]) - ).astype(np.float32) - if not np.isfinite(result).all(): - raise StarVLAError("legacy Qwen2.5 OFT unnormalization produced non-finite actions") - return np.ascontiguousarray(result) - - -@contextlib.contextmanager -def _config_only_qwen25_bootstrap(torch: Any, transformers: Any, qwen_dir: Path): - """Build Qwen2.5 topology without loading a duplicate base weight set.""" - - model_class = transformers.Qwen2_5_VLForConditionalGeneration - had_override = "from_pretrained" in model_class.__dict__ - original_override = model_class.__dict__.get("from_pretrained") - - def from_config_only(model_id: str | os.PathLike[str], **kwargs: Any): - actual = Path(model_id).resolve() - if actual != qwen_dir.resolve(): - raise StarVLAError(f"official wrapper requested unexpected Qwen source: {actual}") - torch_dtype = kwargs.get("torch_dtype") - if torch_dtype not in (None, "auto", torch.bfloat16): - raise StarVLAError(f"unexpected Qwen bootstrap torch_dtype: {torch_dtype!r}") - config = transformers.AutoConfig.from_pretrained( - actual, - local_files_only=True, - trust_remote_code=False, - ) - declared_model_type = getattr(type(config), "model_type", None) - runtime_model_type = getattr(config, "model_type", None) - text_config = getattr(config, "text_config", config) - vision_config = getattr(config, "vision_config", None) - config_contract = { - "declared_model_type": declared_model_type, - "runtime_model_type": runtime_model_type, - "hidden_size": getattr(text_config, "hidden_size", None), - "layer_count": getattr(text_config, "num_hidden_layers", None), - "vocab_size": getattr(text_config, "vocab_size", None), - "vision_hidden_size": getattr(vision_config, "hidden_size", 1280), - "vision_depth": getattr(vision_config, "depth", 32), - "vision_output_size": getattr(vision_config, "out_hidden_size", 2048), - } - expected_contract = { - "declared_model_type": BACKBONE, - # Transformers 4.57 delegates this instance property to text_config. - "runtime_model_type": runtime_model_type, - "hidden_size": 2048, - "layer_count": 36, - "vocab_size": 151936, - "vision_hidden_size": 1280, - "vision_depth": 32, - "vision_output_size": 2048, - } - if ( - runtime_model_type not in {BACKBONE, "qwen2_5_vl_text"} - or config_contract != expected_contract - ): - raise StarVLAError( - f"unexpected local Qwen config contract: {config_contract}" - ) - previous_dtype = torch.get_default_dtype() - try: - torch.set_default_dtype(torch.bfloat16) - with transformers.modeling_utils.no_init_weights(): - model = model_class(config) - finally: - torch.set_default_dtype(previous_dtype) - return model - - model_class.from_pretrained = staticmethod(from_config_only) - try: - yield - finally: - if had_override: - model_class.from_pretrained = original_override - else: - delattr(model_class, "from_pretrained") - - -@contextlib.contextmanager -def _official_qwen25_alias(qwen_dir: Path): - """Give the local processor directory the dispatch name used by StarVLA.""" - - qwen_dir = qwen_dir.resolve() - with tempfile.TemporaryDirectory(prefix="starvla-qwen25-alias-") as temporary: - alias = Path(temporary) / "Qwen2.5-VL-3B-Instruct" - alias.symlink_to(qwen_dir, target_is_directory=True) - if alias.resolve() != qwen_dir: - raise StarVLAError(f"temporary Qwen2.5 alias has the wrong target: {alias}") - yield alias - - -def load_official_framework(paths: Mapping[str, Any], *, device: str) -> tuple[Any, dict[str, Any]]: - import torch - import transformers - - source_dir = Path(paths["source_dir"]) - if any(name == "starVLA" or name.startswith("starVLA.") for name in sys.modules): - raise StarVLAError("starVLA was imported before local source verification") - sys.path.insert(0, str(source_dir)) - try: - from starVLA.model.framework import base_framework, share_tools - from starVLA.model.framework.VLM4A import QwenOFT - - _assert_module_origin(base_framework, source_dir) - _assert_module_origin(share_tools, source_dir) - _assert_module_origin(QwenOFT, source_dir) - config, norm_stats = share_tools.read_mode_config(str(paths["checkpoint"])) - with _official_qwen25_alias(Path(paths["qwen_dir"])) as qwen_alias: - config = base_framework.merge_config_overrides( - config, - [ - f"framework.qwenvl.base_vlm={qwen_alias}", - "framework.qwenvl.attn_implementation=sdpa", - ], - ) - cfg = share_tools.dict_to_namespace(config) - cfg.trainer.pretrained_checkpoint = None - with _config_only_qwen25_bootstrap(torch, transformers, Path(paths["qwen_dir"])): - framework = QwenOFT.Qwenvl_OFT(cfg) - - try: - state_dict = torch.load(paths["checkpoint"], map_location="cpu", weights_only=True) - except TypeError: - state_dict = torch.load(paths["checkpoint"], map_location="cpu") - if not isinstance(state_dict, Mapping) or not state_dict: - raise StarVLAError("official checkpoint did not contain a non-empty state_dict") - framework.load_state_dict(state_dict, strict=True) - del state_dict - gc.collect() - framework.norm_stats = norm_stats - - if type(framework).__name__ != "Qwenvl_OFT": - raise StarVLAError(f"unexpected official framework class: {type(framework).__name__}") - if framework.action_token != ACTION_TOKEN: - raise StarVLAError(f"unexpected OFT action token: {framework.action_token!r}") - token_ids = framework.qwen_vl_interface.processor.tokenizer( - ACTION_TOKEN, add_special_tokens=False - )["input_ids"] - if token_ids != [int(framework.action_token_id)]: - raise StarVLAError(f"Qwen2.5 action token is not one tokenizer token: {token_ids}") - if int(framework.chunk_len) != 16: - raise StarVLAError(f"official OFT action horizon must be 16, got {framework.chunk_len}") - if int(framework.action_model.action_dim) != 7: - raise StarVLAError( - f"official OFT action dimension must be 7, got {framework.action_model.action_dim}" - ) - hidden_size = int(framework.qwen_vl_interface.model.config.hidden_size) - if hidden_size != 2048: - raise StarVLAError(f"official Qwen2.5 hidden size must be 2048, got {hidden_size}") - qwen_dtypes = {parameter.dtype for parameter in framework.qwen_vl_interface.parameters()} - policy_dtypes = {parameter.dtype for parameter in framework.action_model.parameters()} - if qwen_dtypes != {torch.bfloat16}: - raise StarVLAError(f"unexpected Qwen2.5 parameter dtypes: {qwen_dtypes}") - if policy_dtypes != {torch.float32}: - raise StarVLAError(f"unexpected OFT parameter dtypes: {policy_dtypes}") - framework = framework.to(dtype=torch.bfloat16).to(device).eval() - if {parameter.dtype for parameter in framework.parameters()} != {torch.bfloat16}: - raise StarVLAError("official --use_bf16 cast did not cover the whole OFT model") - return framework, config - finally: - if sys.path and sys.path[0] == str(source_dir): - del sys.path[0] - - -def run_official_forward(framework: Any, *, images: Sequence[Any], task: str) -> dict[str, Any]: - """Execute Qwenvl_OFT.predict_action and capture the exact policy boundary.""" - - import torch - - captures: dict[str, Any] = {} - qwen = framework.qwen_vl_interface - action_model = framework.action_model - original_build = qwen.build_qwenvl_inputs - original_gather = framework._gather_action_token_embeddings - original_policy = action_model.predict_action - - def capture_build(*args: Any, **kwargs: Any): - batch_images = kwargs.get("images", args[0] if args else None) - instructions = kwargs.get("instructions", args[1] if len(args) > 1 else None) - captures["processed_images"] = list(batch_images[0]) - captures["framework_instructions"] = list(instructions) - result = original_build(*args, **kwargs) - captures["qwen_inputs"] = { - key: value.detach() - for key, value in result.items() - if isinstance(value, torch.Tensor) - } - return result - - def capture_gather(*args: Any, **kwargs: Any): - queries = original_gather(*args, **kwargs) - captures["action_queries_raw"] = queries.detach() - policy_dtype = next(action_model.parameters()).dtype - captures["policy_input_dtype"] = str(policy_dtype).removeprefix("torch.") - if queries.dtype != policy_dtype: - raise StarVLAError( - f"official whole-model BF16 dtype mismatch: queries={queries.dtype}, policy={policy_dtype}" - ) - return queries - - def capture_policy(*args: Any, **kwargs: Any): - captures["action_queries_policy"] = args[0].detach() - output = original_policy(*args, **kwargs) - captures["raw_policy"] = output.detach() - return output - - def capture_hidden(_module: Any, _inputs: Any, output: Any): - if not getattr(output, "hidden_states", None): - raise StarVLAError("official Qwen2.5 output did not include hidden_states") - captures["last_hidden_state"] = output.hidden_states[-1].detach() - - qwen.build_qwenvl_inputs = capture_build - framework._gather_action_token_embeddings = capture_gather - action_model.predict_action = capture_policy - hook = qwen.register_forward_hook(capture_hidden) - try: - result = framework.predict_action(examples=[{"image": list(images), "lang": task}]) - finally: - hook.remove() - qwen.build_qwenvl_inputs = original_build - framework._gather_action_token_embeddings = original_gather - action_model.predict_action = original_policy - - required = { - "processed_images", - "framework_instructions", - "qwen_inputs", - "action_queries_raw", - "action_queries_policy", - "raw_policy", - "last_hidden_state", - } - missing = sorted(required - set(captures)) - if missing: - raise StarVLAError(f"official Qwen2.5 OFT instrumentation missed: {missing}") - - input_ids, _ = _tensor_to_array(captures["qwen_inputs"]["input_ids"]) - _, selected = select_action_positions( - input_ids, - action_token_id=int(framework.action_token_id), - chunk_len=int(framework.chunk_len), - ) - last_hidden = captures["last_hidden_state"] - positions = torch.as_tensor(selected, device=last_hidden.device, dtype=torch.long) - expected_queries = last_hidden.gather( - 1, positions.unsqueeze(-1).expand(-1, -1, last_hidden.shape[-1]) - ) - if not torch.equal(expected_queries, captures["action_queries_raw"]): - raise StarVLAError("captured action queries do not match Qwen2.5 result_norm positions") - if captures["action_queries_raw"].dtype != torch.bfloat16: - raise StarVLAError( - f"unexpected raw action-query dtype: {captures['action_queries_raw'].dtype}" - ) - if captures["action_queries_policy"].dtype != torch.bfloat16: - raise StarVLAError( - f"unexpected OFT policy input dtype: {captures['action_queries_policy'].dtype}" - ) - if not torch.equal(captures["action_queries_raw"], captures["action_queries_policy"]): - raise StarVLAError("OFT policy input changed across the BF16 model boundary") - - normalized = np.asarray(result.get("normalized_actions")) - raw_policy, _ = _tensor_to_array(captures["raw_policy"]) - expected_shape = (1, int(framework.chunk_len), int(action_model.action_dim)) - if normalized.shape != expected_shape: - raise StarVLAError( - f"official OFT output shape mismatch: expected {expected_shape}, got {normalized.shape}" - ) - if normalized.shape != raw_policy.shape or not np.array_equal(normalized, raw_policy): - raise StarVLAError("normalized_actions differ from the captured OFT policy output") - if not np.isfinite(normalized).all(): - raise StarVLAError("official Qwen2.5 OFT produced non-finite actions") - captures["normalized_actions"] = np.ascontiguousarray(normalized, dtype=np.float32) - return captures - - -def _load_images(image_paths: Iterable[Path]) -> tuple[list[Any], list[dict[str, Any]]]: - try: - from PIL import Image - except ImportError as exc: - raise StarVLAError("Pillow is required to load oracle images") from exc - - images: list[Any] = [] - records: list[dict[str, Any]] = [] - for path in image_paths: - path = path.resolve() - _ensure_regular_file(path, label="input image") - try: - with Image.open(path) as opened: - opened.load() - image = opened.copy() - except (OSError, ValueError) as exc: - raise StarVLAError(f"failed to decode input image {path}: {exc}") from exc - images.append(image) - records.append( - { - "source_path": str(path), - "source_size": path.stat().st_size, - "source_sha256": sha256_file(path), - "decoded_mode": image.mode, - "decoded_size": list(image.size), - "decoded_pixel_sha256": _image_pixel_sha256(image), - } - ) - return images, records - - -def _render_model_prompt(framework: Any, images: Sequence[Any], instruction: str) -> str: - messages = [ - { - "role": "user", - "content": [ - *({"type": "image", "image": image} for image in images), - {"type": "text", "text": instruction}, - ], - } - ] - rendered = framework.qwen_vl_interface.processor.apply_chat_template( - messages, tokenize=False, add_generation_prompt=True - ) - if not isinstance(rendered, str): - raise StarVLAError(f"Qwen2.5 processor returned a non-string prompt: {type(rendered)}") - return rendered - - -def _build_arrays( - captures: Mapping[str, Any], unnormalized: np.ndarray -) -> tuple[dict[str, np.ndarray], dict[str, Any]]: - arrays: dict[str, np.ndarray] = {} - records: dict[str, Any] = {} - - def add(name: str, value: Any) -> None: - if isinstance(value, np.ndarray): - array = np.ascontiguousarray(value) - source_dtype = None - else: - array, source_dtype = _tensor_to_array(value) - arrays[name] = array - records[name] = _array_record(array, source_dtype=source_dtype) - - for key, tensor in sorted(captures["qwen_inputs"].items()): - add(f"qwen_input__{key}", tensor) - add("last_hidden_state", captures["last_hidden_state"]) - add("action_queries_raw", captures["action_queries_raw"]) - add("action_queries_policy", captures["action_queries_policy"]) - add("raw_policy", captures["raw_policy"]) - add("normalized_actions", captures["normalized_actions"]) - add("unnormalized_actions", np.ascontiguousarray(unnormalized, dtype=np.float32)) - return arrays, records - - -def write_golden( - *, - output_dir: Path, - paths: Mapping[str, Any], - framework: Any, - config: Mapping[str, Any], - image_paths: Sequence[Path], - source_image_records: Sequence[Mapping[str, Any]], - task: str, - unnorm_key: str, - captures: Mapping[str, Any], - unnormalized: np.ndarray, -) -> Path: - import torch - import transformers - - output_dir = output_dir.resolve() - if output_dir.exists(): - raise StarVLAError(f"golden output directory already exists: {output_dir}") - output_dir.parent.mkdir(parents=True, exist_ok=True) - - arrays, array_records = _build_arrays(captures, unnormalized) - input_ids = arrays["qwen_input__input_ids"] - if len(image_paths) != 1 or len(source_image_records) != 1: - raise StarVLAError("Qwen2.5 OFT Bridge parity requires exactly one image") - image_grid = arrays.get("qwen_input__image_grid_thw") - pixel_values = arrays.get("qwen_input__pixel_values") - if ( - image_grid is None - or image_grid.shape != (1, 3) - or not np.issubdtype(image_grid.dtype, np.integer) - ): - raise StarVLAError( - "official Qwen2.5 processor must return one image_grid_thw row" - ) - grid_thw = [int(value) for value in image_grid[0]] - if any(value <= 0 for value in grid_thw) or grid_thw[0] != 1: - raise StarVLAError(f"unexpected Qwen2.5 image_grid_thw: {grid_thw}") - patch_count = int(np.prod(grid_thw, dtype=np.int64)) - if ( - pixel_values is None - or pixel_values.ndim != 2 - or pixel_values.shape[0] != patch_count - ): - raise StarVLAError( - "Qwen2.5 pixel_values do not match the processor image grid" - ) - patch_size = 14 - spatial_merge_size = 2 - if patch_count % (spatial_merge_size * spatial_merge_size) != 0: - raise StarVLAError("Qwen2.5 image grid is not divisible by spatial merge size") - resized_size = [grid_thw[2] * patch_size, grid_thw[1] * patch_size] - merged_image_token_count = patch_count // ( - spatial_merge_size * spatial_merge_size - ) - action_token_id = int(framework.action_token_id) - all_positions, selected_positions = select_action_positions( - input_ids, - action_token_id=action_token_id, - chunk_len=int(framework.chunk_len), - ) - framework_instruction = expected_framework_instruction( - config, task, int(framework.chunk_len) - ) - if captures["framework_instructions"] != [framework_instruction]: - raise StarVLAError("official Qwen2.5 OFT prompt construction drifted") - model_instruction = expected_model_instruction(config, framework_instruction) - rendered_prompt = _render_model_prompt( - framework, captures["processed_images"], model_instruction - ) - token_strings = framework.qwen_vl_interface.processor.tokenizer.convert_ids_to_tokens( - input_ids[0].tolist() - ) - - identity = { - "schema_version": SCHEMA_VERSION, - "kind": GOLDEN_KIND, - "checkpoint_sha256": paths["checkpoint_sha256"], - "starvla_revision": paths["source_revision"], - "qwen_assets_sha256": paths["qwen_assets_sha256"], - "task": task, - "unnorm_key": unnorm_key, - "state": [], - "images": [record["source_sha256"] for record in source_image_records], - } - golden_id = _sha256_bytes(_canonical_json(identity)) - - with tempfile.TemporaryDirectory( - prefix=f".{output_dir.name}.", dir=output_dir.parent - ) as temporary: - staging = Path(temporary) - inputs_dir = staging / "inputs" - inputs_dir.mkdir() - image_records: list[dict[str, Any]] = [] - for index, (source_path, source_record) in enumerate( - zip(image_paths, source_image_records, strict=True) - ): - suffix = source_path.suffix.lower() if source_path.suffix else ".img" - artifact = inputs_dir / f"image-{index:02d}{suffix}" - shutil.copyfile(source_path, artifact) - image_records.append( - { - **source_record, - "artifact": artifact.relative_to(staging).as_posix(), - "artifact_size": artifact.stat().st_size, - "artifact_sha256": sha256_file(artifact), - } - ) - - tensors_path = staging / "tensors.npz" - np.savez(tensors_path, **arrays) - manifest = { - "schema_version": SCHEMA_VERSION, - "kind": GOLDEN_KIND, - "golden_id": golden_id, - "created_utc": dt.datetime.now(dt.timezone.utc).isoformat(), - "model_type": MODEL_TYPE, - "backbone": BACKBONE, - "source": { - "checkpoint_repo_id": OFFICIAL_CHECKPOINT_REPO_ID, - "checkpoint_revision": OFFICIAL_CHECKPOINT_REVISION, - "checkpoint_filename": OFFICIAL_CHECKPOINT_FILENAME, - "checkpoint_path": str(paths["checkpoint"]), - "checkpoint_size": paths["checkpoint_size"], - "checkpoint_sha256": paths["checkpoint_sha256"], - "config_yaml_path": str(paths["config_yaml"]), - "config_yaml_size": paths["config_yaml"].stat().st_size, - "config_yaml_sha256": sha256_file(paths["config_yaml"]), - "dataset_statistics_path": str(paths["dataset_statistics"]), - "dataset_statistics_size": paths["dataset_statistics"].stat().st_size, - "dataset_statistics_sha256": sha256_file(paths["dataset_statistics"]), - "qwen_model_path": str(paths["qwen_dir"]), - "qwen_assets": paths["qwen_assets"], - "qwen_assets_sha256": paths["qwen_assets_sha256"], - "starvla_checkout": str(paths["source_dir"]), - "starvla_revision": paths["source_revision"], - }, - "runtime": { - **_runtime_record( - torch, transformers, str(next(framework.parameters()).device) - ), - "qwen-vl-utils": _distribution_version("qwen-vl-utils"), - }, - "determinism": { - "seed": 0, - "torch_deterministic_algorithms": True, - "cublas_workspace_config": os.environ.get("CUBLAS_WORKSPACE_CONFIG"), - "allow_tf32": False, - "attention_implementation": "sdpa", - }, - "input": { - "task": task, - "unnorm_key": unnorm_key, - "state": [], - "images": image_records, - "processed_images": [ - { - "index": index, - "mode": image.mode, - "size": list(image.size), - "pixel_sha256": _image_pixel_sha256(image), - } - for index, image in enumerate(captures["processed_images"]) - ], - }, - "normalization": legacy_normalization_contract( - paths["norm_stats"], unnorm_key - ), - "model_contract": { - "framework_class": f"{type(framework).__module__}.{type(framework).__name__}", - "action_token": ACTION_TOKEN, - "action_token_id": action_token_id, - "action_horizon": int(framework.chunk_len), - "action_dim": int(framework.action_model.action_dim), - "qwen_hidden_dim": int(framework.qwen_vl_interface.model.config.hidden_size), - "policy_input_dtype": captures["policy_input_dtype"], - "processor_class": type( - framework.qwen_vl_interface.processor - ).__name__, - "image_processor_class": type( - framework.qwen_vl_interface.processor.image_processor - ).__name__, - "image_patch_size": patch_size, - "image_spatial_merge_size": spatial_merge_size, - "image_grid_thw": grid_thw, - "image_resized_size": resized_size, - "merged_image_token_count": merged_image_token_count, - }, - "prompt": { - "framework_instruction": framework_instruction, - "model_instruction": model_instruction, - "rendered_chat_template": rendered_prompt, - }, - "tokens": { - "input_ids": input_ids.tolist(), - "token_strings": token_strings, - "all_action_token_positions": all_positions, - "selected_action_token_positions": selected_positions, - }, - "outputs": { - "normalized_actions": arrays["normalized_actions"].tolist(), - "unnormalized_actions": arrays["unnormalized_actions"].tolist(), - }, - "action_gate": { - "metric": "full_tensor_global_relative_l2", - "operator": "<=", - "limit": ACTION_RELATIVE_L2_LIMIT, - "required_outputs": ["normalized_actions", "unnormalized_actions"], - }, - "artifacts": { - "tensors": { - "path": tensors_path.name, - "size": tensors_path.stat().st_size, - "sha256": sha256_file(tensors_path), - "arrays": array_records, - } - }, - } - manifest_path = staging / "golden.json" - manifest_path.write_text( - json.dumps(manifest, ensure_ascii=False, indent=2, sort_keys=True) + "\n", - encoding="utf-8", - ) - Path(temporary).replace(output_dir) - return output_dir / "golden.json" - - -def build_parser() -> argparse.ArgumentParser: - parser = argparse.ArgumentParser( - description=( - "Generate an auditable local-Python golden from the official " - "StarVLA Qwen2.5-VL OFT .pt checkpoint." - ) - ) - parser.add_argument("--checkpoint", required=True, type=Path) - parser.add_argument( - "--qwen-model", - required=True, - type=Path, - help="Local Qwen2.5-VL-3B-Instruct topology and processor directory", - ) - parser.add_argument( - "--starvla-source", - type=Path, - default=Path("ckpts/starvla/source/starvla"), - ) - parser.add_argument( - "--expected-checkpoint-sha256", - default=OFFICIAL_CHECKPOINT_SHA256, - ) - parser.add_argument( - "--expected-checkpoint-size", - default=OFFICIAL_CHECKPOINT_SIZE, - type=int, - ) - parser.add_argument("--expected-source-revision") - parser.add_argument("--image", action="append", default=[], type=Path) - parser.add_argument("--task") - parser.add_argument("--unnorm-key", choices=tuple(LEGACY_UNNORM_PROFILES)) - parser.add_argument("--output-dir", type=Path) - parser.add_argument("--device", default="cuda:0") - parser.add_argument("--preflight-only", action="store_true") - return parser - - -def main(argv: Sequence[str] | None = None) -> int: - args = build_parser().parse_args(argv) - _require_isolated_python() - if not args.preflight_only: - missing = [ - name - for name, value in ( - ("--image", args.image), - ("--task", args.task), - ("--unnorm-key", args.unnorm_key), - ("--output-dir", args.output_dir), - ) - if not value - ] - if missing: - raise StarVLAError("golden generation requires " + ", ".join(missing)) - if not args.task.strip() or not args.unnorm_key.strip(): - raise StarVLAError("--task and --unnorm-key must not be empty") - if len(args.image) != 1: - raise StarVLAError( - "Qwen2.5 OFT Bridge golden generation requires exactly one --image" - ) - - paths = validate_local_inputs( - checkpoint=args.checkpoint, - qwen_model=args.qwen_model, - source_dir=args.starvla_source, - expected_checkpoint_sha256=args.expected_checkpoint_sha256, - expected_checkpoint_size=args.expected_checkpoint_size, - expected_source_revision=args.expected_source_revision, - ) - try: - import torch - import transformers - except ImportError as exc: - raise StarVLAError(f"official StarVLA runtime dependency is missing: {exc}") from exc - validate_runtime_versions( - torch_version=torch.__version__, - torchvision_version=_distribution_version("torchvision"), - transformers_version=transformers.__version__, - numpy_version=np.__version__, - ) - qwen_vl_utils_version = _distribution_version("qwen-vl-utils") - if qwen_vl_utils_version != EXPECTED_QWEN_VL_UTILS_VERSION: - raise StarVLAError( - "official Qwen2.5 OFT oracle requires qwen-vl-utils " - f"{EXPECTED_QWEN_VL_UTILS_VERSION}, got {qwen_vl_utils_version}" - ) - _configure_determinism(torch, seed=0, device=args.device) - if args.preflight_only: - print( - "Qwen2.5 OFT local .pt preflight passed: " - f"{paths['checkpoint']} ({paths['checkpoint_sha256']})" - ) - return 0 - - images, image_records = _load_images(args.image) - framework, config = load_official_framework(paths, device=args.device) - captures = run_official_forward(framework, images=images, task=args.task) - - normalized = captures["normalized_actions"] - unnormalized = unnormalize_legacy_actions( - normalized, paths["norm_stats"], args.unnorm_key - ) - - manifest = write_golden( - output_dir=args.output_dir, - paths=paths, - framework=framework, - config=config, - image_paths=args.image, - source_image_records=image_records, - task=args.task, - unnorm_key=args.unnorm_key, - captures=captures, - unnormalized=np.ascontiguousarray(unnormalized, dtype=np.float32), - ) - print(f"Wrote StarVLA Qwen2.5 OFT local-Python golden: {manifest}") - return 0 - - -if __name__ == "__main__": - try: - raise SystemExit(main()) - except StarVLAError as exc: - raise SystemExit(f"error: {exc}") from exc diff --git a/tools/hf2gguf/starvla/generate_starvla_qwen25_pi_golden.py b/tools/hf2gguf/starvla/generate_starvla_qwen25_pi_golden.py deleted file mode 100644 index b83e915..0000000 --- a/tools/hf2gguf/starvla/generate_starvla_qwen25_pi_golden.py +++ /dev/null @@ -1,1338 +0,0 @@ -#!/usr/bin/env python3 -"""Generate a fixed-noise local-Python oracle for the released Qwen2.5 PI. - -The published checkpoint predates the current QwenPI refactor. This exporter -therefore executes the exact historical implementation stored in the pinned -local StarVLA git repository, including the documented ``--use_bf16`` -deployment path. It applies one bootstrap shim: - -* construct Qwen2.5-VL from its local config before the complete checkpoint is - loaded, avoiding a duplicate base-weight download. - -After strict loading, the whole framework is converted to BF16 exactly as in -the official server command. The action head remains the historical -16-block, all-cross-attention forward. -Its initial 16x7 noise tensor is an explicit binary-fraction fixture shared -with the C++ parity runner; cross-language RNG replay is never used. -""" - -from __future__ import annotations - -import argparse -import contextlib -import datetime as dt -import gc -import hashlib -import io -import json -import math -import os -import shutil -import subprocess -import sys -import tarfile -import tempfile -from pathlib import Path -from typing import Any, Iterable, Mapping, Sequence - -import numpy as np - - -TOOLS_DIR = Path(__file__).resolve().parent -if str(TOOLS_DIR) not in sys.path: - sys.path.insert(0, str(TOOLS_DIR)) - -from generate_starvla_oft_golden import ( # noqa: E402 - _assert_module_origin, - _canonical_json, - _configure_determinism, - _distribution_version, - _ensure_regular_file, - _image_pixel_sha256, - _require_isolated_python, - _runtime_record, - _sha256_bytes, - validate_runtime_versions, -) -from generate_starvla_qwen25_groot_golden import ( # noqa: E402 - ACTION_TOKEN_COUNT, - ACTION_TOKEN_ID_MAX, - ACTION_TOKEN_ID_MIN, - EXPECTED_QWEN_VL_UTILS_VERSION, - validate_action_tokenizer_assets, - validate_processor_contract, -) -from generate_starvla_qwen25_oft_golden import ( # noqa: E402 - _official_qwen25_alias, - _verify_clean_source, -) -from starvla_checkpoint import ( # noqa: E402 - DEFAULT_CATALOG, - StarVLAError, - get_variant, - load_catalog, - official_bundle_uuid, - sha256_file, - verify_catalog_files, - verify_checkpoint_file, -) - - -SCHEMA_VERSION = 1 -GOLDEN_KIND = "starvla_qwen25_pi_local_pt_python_oracle" -MODEL_TYPE = "starvla" -VARIANT = "qwen25_pi" -BACKBONE = "qwen2_5_vl" -ACTION_RELATIVE_L2_LIMIT = 0.03 - -OFFICIAL_CHECKPOINT_REPO_ID = "StarVLA/Qwen-PI-Bridge-RT-1" -OFFICIAL_CHECKPOINT_REVISION = "26d0e079fbe3bc3fc62301f44f0025ef7c64ee22" -OFFICIAL_CHECKPOINT_FILENAME = "steps_30000_pytorch_model.pt" -OFFICIAL_CHECKPOINT_SIZE = 10_103_104_403 -OFFICIAL_CHECKPOINT_SHA256 = ( - "8a0e47858921924d5038f7c4393dee6682b83175a85546e35e357e8f74ce8343" -) -OFFICIAL_QWEN_REPO_ID = "StarVLA/Qwen2.5-VL-3B-Instruct-Action" -OFFICIAL_QWEN_REVISION = "ce86bd9a53416527b8361e8dfc47316288ffa110" -OFFICIAL_STARVLA_REPO_ID = "starVLA/starVLA" -OFFICIAL_STARVLA_REVISION = "631aae02afe6d95876e923ff518e8ff2ab9a2f88" -LEGACY_IMPLEMENTATION_REVISION = "e872a8579055f9332add8a2549b9fd5599e11510" -PI_RUNTIME_CONTRACT_SHA256 = ( - "dea02dbb9099b34454db473c39375ce6467109287a27d4ab89193561be035219" -) - -EXPECTED_ACTION_HORIZON = 16 -EXPECTED_ACTION_DIM = 7 -EXPECTED_STATE_DIM = 7 -EXPECTED_QWEN_HIDDEN_DIM = 2048 -EXPECTED_QWEN_LAYER_COUNT = 36 -EXPECTED_HIDDEN_TUPLE_INDICES = list(range(21, 37)) -EXPECTED_DIT_BLOCK_COUNT = 16 -EXPECTED_DIT_WIDTH = 2048 -EXPECTED_FUTURE_TOKEN_COUNT = 32 -EXPECTED_TIMESTEP_IDS = [0, 250, 500, 750] -EXPECTED_COT_TEMPLATE = ( - "Your task is {instruction}. To identify the key objects for your task. " - "Locate their bounding boxes in [x1,y1,x2,y2] format." -) -UNNORM_KEYS = ("oxe_bridge", "oxe_rt1") - -NOISE_ALGORITHM = "portable_binary_fraction_lcg_v1" -NOISE_DENOMINATOR = 64 -NOISE_MULTIPLIER = 73 -NOISE_INCREMENT = 19 -NOISE_MODULUS = 257 -NOISE_OFFSET = 128 - -LEGACY_SOURCE_FILES = { - "deployment/model_server/README.md": - "85662206d8f9ba1948ccc2c588b241fe9f45e0ee43e77ce825a2247f195cb3a6", - "deployment/model_server/server_policy.py": - "98569a4d3a1781d9c9b0fa5bd1952c4f212ff5307522078544bee59059a7df17", - "examples/LIBERO/model2libero_interface.py": - "16f760d011513f6be4f6fbc304aa6567f0a517015122f8c7f67ac85a07f37a13", - "examples/SimplerEnv/model2simpler_interface.py": - "510ae919871ccd6ce64271338c9dfb01648a3c0a820adf90998537ed9bf0fac3", - "starVLA/model/framework/base_framework.py": - "12cdfc8afbff72a44e3f4d0bbafc229721e49de26fe97c5b79821db06118c334", - "starVLA/model/framework/QwenPI.py": - "d368c669ec178045ca4143c7f90c6db75082946042afe915d4686e18d43be525", - "starVLA/model/modules/action_model/LayerwiseFM_ActionHeader.py": - "c586021a5d98605c01728d3ccc98218ce3bc639a95c06f1579eb8289612f5d43", - "starVLA/model/modules/action_model/flow_matching_head/cross_attention_dit.py": - "d835796a351f4562b826ada959332c7baa063b79432ca15e4d0ce76745128a62", - "starVLA/model/modules/vlm/QWen2_5.py": - "b94b9220a04ad6017789e9e907fc2d6e7ec8c32f77a7a38d4adebd18bf2fe5c3", -} - - -def explicit_initial_noise() -> np.ndarray: - """Return the portable 16x7 parity noise using exact binary fractions.""" - - count = EXPECTED_ACTION_HORIZON * EXPECTED_ACTION_DIM - index = np.arange(count, dtype=np.int64) - numerator = ( - (index * NOISE_MULTIPLIER + NOISE_INCREMENT) % NOISE_MODULUS - ) - NOISE_OFFSET - result = numerator.astype(np.float32) / np.float32(NOISE_DENOMINATOR) - return np.ascontiguousarray( - result.reshape(1, EXPECTED_ACTION_HORIZON, EXPECTED_ACTION_DIM) - ) - - -def _load_json_object(path: Path, *, label: str) -> dict[str, Any]: - _ensure_regular_file(path, label=label) - try: - value = json.loads(path.read_text(encoding="utf-8")) - except (OSError, UnicodeError, json.JSONDecodeError) as exc: - raise StarVLAError(f"failed to parse {label} {path}: {exc}") from exc - if not isinstance(value, dict): - raise StarVLAError(f"{label} root must be an object") - return value - - -def _run_git(source_dir: Path, *args: str, binary: bool = False) -> bytes | str: - try: - result = subprocess.run( - ["git", "-C", str(source_dir), *args], - check=True, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - ) - except (OSError, subprocess.CalledProcessError) as exc: - detail = "" - if isinstance(exc, subprocess.CalledProcessError): - detail = exc.stderr.decode("utf-8", errors="replace").strip() - raise StarVLAError( - f"failed to inspect pinned StarVLA git object: {detail or exc}" - ) from exc - return result.stdout if binary else result.stdout.decode("utf-8").strip() - - -def verify_source_semantics(source_dir: Path) -> dict[str, Any]: - """Bind the exact legacy implementation used by the released checkpoint.""" - - object_type = _run_git( - source_dir, "cat-file", "-t", LEGACY_IMPLEMENTATION_REVISION - ) - if object_type != "commit": - raise StarVLAError( - f"legacy PI revision is not a commit: {LEGACY_IMPLEMENTATION_REVISION}" - ) - - files: dict[str, str] = {} - sources: dict[str, str] = {} - for relative, expected in LEGACY_SOURCE_FILES.items(): - payload = _run_git( - source_dir, - "show", - f"{LEGACY_IMPLEMENTATION_REVISION}:{relative}", - binary=True, - ) - assert isinstance(payload, bytes) - digest = hashlib.sha256(payload).hexdigest() - if digest != expected: - raise StarVLAError( - f"legacy PI source SHA256 mismatch for {relative}: " - f"expected {expected}, got {digest}" - ) - files[relative] = digest - sources[relative] = payload.decode("utf-8") - - framework_source = sources["starVLA/model/framework/QwenPI.py"] - action_source = sources[ - "starVLA/model/modules/action_model/LayerwiseFM_ActionHeader.py" - ] - base_source = sources["starVLA/model/framework/base_framework.py"] - server_source = sources["deployment/model_server/server_policy.py"] - deployment_readme = sources["deployment/model_server/README.md"] - evaluator_sources = ( - sources["examples/LIBERO/model2libero_interface.py"], - sources["examples/SimplerEnv/model2simpler_interface.py"], - ) - required_framework = ( - "expected_layers = len(self.action_model.model.transformer_blocks)", - "vl_embs_list = list(all_hidden[-expected_layers:])", - 'getattr(self.config.datasets.vla_data, "image_size", None)', - 'with torch.autocast("cuda", dtype=torch.float32):', - ) - required_action = ( - "for layer_idx, layer in enumerate(self.model.transformer_blocks):", - "encoder_hidden_states=vl_embs_list[layer_idx]", - "actions = torch.randn(", - "actions = actions + dt * pred_velocity", - ) - required_normalization = ( - "normalized_actions = np.clip(normalized_actions, -1, 1)", - "normalized_actions[:, 6] = np.where(normalized_actions[:, 6] < 0.5, 0, 1)", - ) - missing = [ - fragment - for fragment in required_framework - if fragment not in framework_source - ] + [ - fragment for fragment in required_action if fragment not in action_source - ] + [ - fragment for fragment in required_normalization if fragment not in base_source - ] + [ - fragment - for fragment in ("vla = vla.to(torch.bfloat16)",) - if fragment not in server_source - ] + [ - fragment - for fragment in ("--use_bf16",) - if fragment not in deployment_readme - ] + [ - fragment - for evaluator_source in evaluator_sources - for fragment in required_normalization - if fragment not in evaluator_source - ] - if missing: - raise StarVLAError( - f"legacy PI source semantics probe failed: {missing!r}" - ) - return { - "revision": LEGACY_IMPLEMENTATION_REVISION, - "files": files, - "checkpoint_block_count": EXPECTED_DIT_BLOCK_COUNT, - "hidden_tuple_indices": EXPECTED_HIDDEN_TUPLE_INDICES, - "block_mode": "layerwise_cross_attention_every_block", - "released_config_interleave_self_attention": True, - "use_canonical_dit_forward": False, - "attention_mask_runtime_active": False, - "deployment_precision": "whole_model_bf16_via_use_bf16", - "normalization": "clip_minus1_plus1_then_binary_ge_0_5", - } - - -def _validate_catalog_identity( - catalog: Mapping[str, Any], -) -> tuple[dict[str, Any], dict[str, Any]]: - variant = get_variant(catalog, VARIANT) - qwen_key = variant.get("qwen_asset") - qwen = catalog.get("shared_assets", {}).get(qwen_key) - if not isinstance(qwen, dict): - raise StarVLAError(f"catalog variant {VARIANT} has no Qwen action asset") - expected_checkpoint = { - "path": f"checkpoints/{OFFICIAL_CHECKPOINT_FILENAME}", - "size": OFFICIAL_CHECKPOINT_SIZE, - "sha256": OFFICIAL_CHECKPOINT_SHA256, - } - if ( - variant.get("repo_id") != OFFICIAL_CHECKPOINT_REPO_ID - or variant.get("revision") != OFFICIAL_CHECKPOINT_REVISION - or variant.get("checkpoint") != expected_checkpoint - ): - raise StarVLAError("catalog Qwen2.5 PI checkpoint identity drifted") - if ( - qwen.get("repo_id") != OFFICIAL_QWEN_REPO_ID - or qwen.get("revision") != OFFICIAL_QWEN_REVISION - ): - raise StarVLAError("catalog Qwen2.5 PI action-tokenizer identity drifted") - if ( - catalog.get("source_revisions", {}).get("starvla") - != OFFICIAL_STARVLA_REVISION - ): - raise StarVLAError("catalog StarVLA source revision drifted") - return variant, qwen - - -def _validate_effective_config(config: Mapping[str, Any]) -> None: - try: - framework = config["framework"] - action = framework["action_model"] - diffusion = action["diffusion_model_cfg"] - vla = config["datasets"]["vla_data"] - except (KeyError, TypeError) as exc: - raise StarVLAError("effective Qwen2.5 PI config is incomplete") from exc - actual = { - "framework": framework.get("name"), - "base_vlm": framework.get("qwenvl", {}).get("base_vlm"), - "action_model_type": action.get("action_model_type"), - "configured_hidden_size": action.get("hidden_size"), - "historical_runtime_hidden_size": action.get("action_hidden_dim"), - "action_horizon": action.get("action_horizon"), - "future_action_window_size": action.get("future_action_window_size"), - "action_dim": action.get("action_dim"), - "state_dim": action.get("state_dim"), - "steps": action.get("num_inference_timesteps"), - "buckets": action.get("num_timestep_buckets"), - "future_tokens": action.get("num_target_vision_tokens"), - "layers": diffusion.get("num_layers"), - "cross_dim": diffusion.get("cross_attention_dim"), - "output_dim": diffusion.get("output_dim"), - "interleave": diffusion.get("interleave_self_attention"), - "image_size": vla.get("image_size"), - "default_image_resolution": vla.get("default_image_resolution"), - "obs": vla.get("obs"), - "data_mix": vla.get("data_mix"), - "cot": vla.get("CoT_prompt"), - } - expected = { - "framework": "QwenPI", - "base_vlm": "starVLA/Qwen2.5-VL-3B-Instruct-Action", - "action_model_type": "DiT-Qwen", - "configured_hidden_size": 1024, - "historical_runtime_hidden_size": EXPECTED_DIT_WIDTH, - "action_horizon": EXPECTED_ACTION_HORIZON, - "future_action_window_size": EXPECTED_ACTION_HORIZON - 1, - "action_dim": EXPECTED_ACTION_DIM, - "state_dim": EXPECTED_STATE_DIM, - "steps": 4, - "buckets": 1000, - "future_tokens": EXPECTED_FUTURE_TOKEN_COUNT, - "layers": EXPECTED_DIT_BLOCK_COUNT, - "cross_dim": EXPECTED_QWEN_HIDDEN_DIM, - "output_dim": 1024, - "interleave": True, - "image_size": [224, 224], - "default_image_resolution": [3, 224, 224], - "obs": ["image_0"], - "data_mix": "bridge_rt_1", - "cot": EXPECTED_COT_TEMPLATE, - } - if actual != expected: - raise StarVLAError(f"unexpected effective Qwen2.5 PI config: {actual}") - - -def normalization_contract( - norm_stats: Mapping[str, Any], unnorm_key: str -) -> dict[str, Any]: - if unnorm_key not in UNNORM_KEYS: - raise StarVLAError( - f"Qwen2.5 PI unnorm_key must be one of {list(UNNORM_KEYS)}" - ) - profile = norm_stats.get(unnorm_key) - action = profile.get("action") if isinstance(profile, Mapping) else None - state = profile.get("state") if isinstance(profile, Mapping) else None - if not isinstance(action, Mapping) or not isinstance(state, Mapping): - raise StarVLAError( - f"dataset statistics has no complete {unnorm_key} action/state objects" - ) - try: - q01 = np.asarray(action["q01"], dtype=np.float32) - q99 = np.asarray(action["q99"], dtype=np.float32) - mask = np.asarray(action["mask"], dtype=np.bool_) - state_q01 = np.asarray(state["q01"], dtype=np.float32) - except (KeyError, TypeError, ValueError) as exc: - raise StarVLAError( - f"invalid Qwen2.5 PI statistics for {unnorm_key}: {exc}" - ) from exc - if q01.shape != (7,) or q99.shape != (7,) or mask.shape != (7,): - raise StarVLAError("Qwen2.5 PI action statistics must be 7D") - if state_q01.shape != (8,): - raise StarVLAError( - "Qwen2.5 PI dataset state statistics must remain 8D" - ) - if ( - not np.isfinite(q01).all() - or not np.isfinite(q99).all() - or np.any(q99[mask] <= q01[mask]) - or mask.tolist() != [True, True, True, True, True, True, False] - ): - raise StarVLAError("Qwen2.5 PI normalization statistics are invalid") - return { - "stats_key": unnorm_key, - "q01": q01.tolist(), - "q99": q99.tolist(), - "mask": mask.tolist(), - "continuous_dimensions": [0, 1, 2, 3, 4, 5], - "binary_dimensions": [6], - "binary_threshold": 0.5, - "binary_comparison": "ge", - "continuous_clip": True, - "state_input_contract": - "caller_supplies_model_7d_state_8d_dataset_stats_are_not_applied", - } - - -def unnormalize_actions( - normalized: np.ndarray, - norm_stats: Mapping[str, Any], - unnorm_key: str, -) -> np.ndarray: - contract = normalization_contract(norm_stats, unnorm_key) - values = np.clip( - np.ascontiguousarray(normalized, dtype=np.float32), - np.float32(-1.0), - np.float32(1.0), - ) - if values.shape != (1, EXPECTED_ACTION_HORIZON, EXPECTED_ACTION_DIM): - raise StarVLAError( - f"normalized Qwen2.5 PI actions have invalid shape: {values.shape}" - ) - q01 = np.asarray(contract["q01"], dtype=np.float32) - q99 = np.asarray(contract["q99"], dtype=np.float32) - mask = np.asarray(contract["mask"], dtype=np.bool_) - output = np.empty_like(values) - output[..., mask] = ( - (values[..., mask] + np.float32(1.0)) - * np.float32(0.5) - * (q99[mask] - q01[mask]) - + q01[mask] - ) - output[..., ~mask] = ( - values[..., ~mask] >= np.float32(contract["binary_threshold"]) - ).astype(np.float32) - if not np.isfinite(output).all(): - raise StarVLAError("Qwen2.5 PI unnormalization produced non-finite values") - return np.ascontiguousarray(output) - - -def validate_local_inputs( - *, - checkpoint_root: Path, - checkpoint: Path | None, - qwen_model: Path | None, - source_dir: Path, - catalog_path: Path = DEFAULT_CATALOG, -) -> dict[str, Any]: - catalog = load_catalog(catalog_path) - variant, qwen = _validate_catalog_identity(catalog) - checkpoint_root = checkpoint_root.resolve() - policy_dir = checkpoint_root / "sources" / variant["directory"] / variant["revision"] - qwen_dir = checkpoint_root / "sources" / qwen["directory"] / qwen["revision"] - checkpoint_path = policy_dir / variant["checkpoint"]["path"] - if checkpoint is not None and checkpoint.resolve() != checkpoint_path.resolve(): - raise StarVLAError( - f"Qwen2.5 PI checkpoint must be the catalog path {checkpoint_path}" - ) - if qwen_model is not None and qwen_model.resolve() != qwen_dir.resolve(): - raise StarVLAError( - f"Qwen2.5 PI processor must be the catalog path {qwen_dir}" - ) - - verify_catalog_files(policy_dir, variant) - verify_catalog_files(qwen_dir, qwen) - tokenizer = validate_action_tokenizer_assets(qwen_dir) - source_dir = source_dir.resolve() - revision = _verify_clean_source(source_dir, OFFICIAL_STARVLA_REVISION) - source_probe = verify_source_semantics(source_dir) - - sidecar = Path(f"{checkpoint_path}.aria2") - checkpoint_ready = ( - checkpoint_path.is_file() - and not checkpoint_path.is_symlink() - and not sidecar.exists() - ) - if checkpoint_ready: - verify_checkpoint_file(checkpoint_path, variant) - - config_yaml = policy_dir / "config.yaml" - dataset_statistics = policy_dir / "dataset_statistics.json" - try: - import yaml - - config = yaml.safe_load(config_yaml.read_text(encoding="utf-8")) - except (ImportError, OSError, UnicodeError, ValueError) as exc: - raise StarVLAError(f"failed to load Qwen2.5 PI config.yaml: {exc}") from exc - if not isinstance(config, dict): - raise StarVLAError("Qwen2.5 PI config.yaml root must be an object") - _validate_effective_config(config) - norm_stats = _load_json_object( - dataset_statistics, label="Qwen2.5 PI dataset statistics" - ) - if set(norm_stats) != set(UNNORM_KEYS): - raise StarVLAError( - f"unexpected Qwen2.5 PI normalization profiles: {sorted(norm_stats)}" - ) - for key in UNNORM_KEYS: - normalization_contract(norm_stats, key) - return { - "catalog": catalog, - "catalog_path": catalog_path.resolve(), - "variant": variant, - "qwen": qwen, - "policy_dir": policy_dir.resolve(), - "qwen_dir": qwen_dir.resolve(), - "checkpoint": checkpoint_path.resolve(), - "checkpoint_ready": checkpoint_ready, - "config_yaml": config_yaml.resolve(), - "dataset_statistics": dataset_statistics.resolve(), - "norm_stats": norm_stats, - "config": config, - "source_dir": source_dir, - "source_revision": revision, - "source_probe": source_probe, - "tokenizer": tokenizer, - } - - -def _extract_legacy_git_archive(archive: bytes, destination: Path) -> None: - """Extract verified Git files while ignoring repository-local links.""" - - def regular_file_filter( - member: tarfile.TarInfo, target: str - ) -> tarfile.TarInfo | None: - # The historical tree contains dataset links to machine-local absolute - # paths. They are irrelevant to inference and must never be followed or - # materialized while constructing the verified runtime source tree. - if member.issym() or member.islnk(): - return None - return tarfile.data_filter(member, target) - - with tarfile.open(fileobj=io.BytesIO(archive), mode="r:") as stream: - stream.extractall(destination, filter=regular_file_filter) - - -@contextlib.contextmanager -def _legacy_source_checkout(source_dir: Path): - """Extract the exact historical source tree without mutating git worktrees.""" - - archive = _run_git( - source_dir, - "archive", - "--format=tar", - LEGACY_IMPLEMENTATION_REVISION, - binary=True, - ) - assert isinstance(archive, bytes) - with tempfile.TemporaryDirectory(prefix="starvla-qwen25-pi-legacy-") as temporary: - root = Path(temporary) - _extract_legacy_git_archive(archive, root) - for relative, expected in LEGACY_SOURCE_FILES.items(): - path = root / relative - _ensure_regular_file(path, label=f"extracted legacy source {relative}") - if sha256_file(path) != expected: - raise StarVLAError( - f"extracted legacy source changed unexpectedly: {relative}" - ) - yield root - - -@contextlib.contextmanager -def _config_only_qwen25_bootstrap( - torch: Any, transformers: Any, qwen_dir: Path -): - """Build the Qwen topology in BF16 without loading absent base weights.""" - - model_class = transformers.Qwen2_5_VLForConditionalGeneration - had_override = "from_pretrained" in model_class.__dict__ - original_override = model_class.__dict__.get("from_pretrained") - - def from_config_only(model_id: str | os.PathLike[str], **kwargs: Any): - actual = Path(model_id).resolve() - if actual != qwen_dir.resolve(): - raise StarVLAError( - f"legacy PI wrapper requested unexpected Qwen source: {actual}" - ) - if kwargs.get("torch_dtype") not in (None, "auto", torch.bfloat16): - raise StarVLAError( - f"unexpected Qwen bootstrap dtype: {kwargs.get('torch_dtype')!r}" - ) - config = transformers.AutoConfig.from_pretrained( - actual, local_files_only=True, trust_remote_code=False - ) - declared_model_type = getattr(type(config), "model_type", None) - runtime_model_type = getattr(config, "model_type", None) - text_config = getattr(config, "text_config", config) - vision_config = getattr(config, "vision_config", None) - config_contract = { - "declared_model_type": declared_model_type, - "runtime_model_type": runtime_model_type, - "hidden_size": getattr(text_config, "hidden_size", None), - "layer_count": getattr(text_config, "num_hidden_layers", None), - "vocab_size": getattr(text_config, "vocab_size", None), - "vision_hidden_size": getattr(vision_config, "hidden_size", 1280), - "vision_depth": getattr(vision_config, "depth", 32), - "vision_output_size": getattr(vision_config, "out_hidden_size", 2048), - } - expected_contract = { - "declared_model_type": BACKBONE, - "runtime_model_type": runtime_model_type, - "hidden_size": 2048, - "layer_count": 36, - "vocab_size": 153713, - "vision_hidden_size": 1280, - "vision_depth": 32, - "vision_output_size": 2048, - } - if ( - runtime_model_type not in {BACKBONE, "qwen2_5_vl_text"} - or config_contract != expected_contract - ): - raise StarVLAError( - f"unexpected local Qwen config contract: {config_contract}" - ) - config._attn_implementation = "sdpa" - previous_dtype = torch.get_default_dtype() - try: - torch.set_default_dtype(torch.bfloat16) - with transformers.modeling_utils.no_init_weights(): - return model_class(config) - finally: - torch.set_default_dtype(previous_dtype) - - model_class.from_pretrained = staticmethod(from_config_only) - try: - yield - finally: - if had_override: - model_class.from_pretrained = original_override - else: - delattr(model_class, "from_pretrained") - - -def load_official_framework( - paths: Mapping[str, Any], *, device: str -) -> tuple[Any, dict[str, Any], tempfile.TemporaryDirectory[str]]: - """Load the original checkpoint against its exact historical Python code.""" - - import torch - import transformers - - if not paths["checkpoint_ready"]: - raise StarVLAError( - f"official Qwen2.5 PI checkpoint is absent or incomplete: " - f"{paths['checkpoint']}" - ) - if any(name == "starVLA" or name.startswith("starVLA.") for name in sys.modules): - raise StarVLAError("starVLA was imported before legacy-source verification") - - # Keep the extracted tree alive for as long as the framework class exists. - holder: tempfile.TemporaryDirectory[str] = tempfile.TemporaryDirectory( - prefix="starvla-qwen25-pi-runtime-" - ) - runtime_root = Path(holder.name) - archive = _run_git( - Path(paths["source_dir"]), - "archive", - "--format=tar", - LEGACY_IMPLEMENTATION_REVISION, - binary=True, - ) - assert isinstance(archive, bytes) - _extract_legacy_git_archive(archive, runtime_root) - - sys.path.insert(0, str(runtime_root)) - try: - from starVLA.model.framework import base_framework, share_tools - from starVLA.model.framework import QwenPI - - for module in (base_framework, share_tools, QwenPI): - _assert_module_origin(module, runtime_root) - config = json.loads(json.dumps(paths["config"])) - with _official_qwen25_alias(Path(paths["qwen_dir"])) as qwen_alias: - config["framework"]["qwenvl"]["base_vlm"] = str(qwen_alias) - cfg = share_tools.dict_to_namespace(config) - cfg.trainer.pretrained_checkpoint = None - with _config_only_qwen25_bootstrap( - torch, transformers, Path(paths["qwen_dir"]) - ): - framework = QwenPI.Qwen_PI(cfg) - - try: - state = torch.load( - paths["checkpoint"], - map_location="cpu", - mmap=True, - weights_only=True, - ) - except TypeError: - state = torch.load( - paths["checkpoint"], map_location="cpu", weights_only=True - ) - if not isinstance(state, Mapping) or not state: - raise StarVLAError("official Qwen2.5 PI checkpoint has no state_dict") - framework.load_state_dict(state, strict=True) - del state - gc.collect() - framework.norm_stats = paths["norm_stats"] - - action_model = framework.action_model - if type(framework).__name__ != "Qwen_PI": - raise StarVLAError( - f"unexpected legacy framework class: {type(framework).__name__}" - ) - if len(action_model.model.transformer_blocks) != EXPECTED_DIT_BLOCK_COUNT: - raise StarVLAError( - "legacy PI checkpoint did not instantiate exactly 16 DiT blocks" - ) - if int(action_model.action_horizon) != EXPECTED_ACTION_HORIZON: - raise StarVLAError("legacy PI action horizon changed") - qwen_dtypes = { - parameter.dtype - for parameter in framework.qwen_vl_interface.parameters() - } - policy_dtypes = { - parameter.dtype for parameter in action_model.parameters() - } - if qwen_dtypes != {torch.bfloat16} or policy_dtypes != {torch.float32}: - raise StarVLAError( - "legacy PI dtype boundary changed: " - f"qwen={qwen_dtypes}, policy={policy_dtypes}" - ) - framework = framework.to(device=device, dtype=torch.bfloat16).eval() - runtime_dtypes = {parameter.dtype for parameter in framework.parameters()} - if runtime_dtypes != {torch.bfloat16}: - raise StarVLAError( - f"official --use_bf16 deployment cast failed: {runtime_dtypes}" - ) - return framework, config, holder - except Exception: - holder.cleanup() - raise - finally: - if sys.path and sys.path[0] == str(runtime_root): - del sys.path[0] - - -def _load_images( - image_paths: Iterable[Path], -) -> tuple[list[Any], list[dict[str, Any]]]: - from PIL import Image - - images: list[Any] = [] - records: list[dict[str, Any]] = [] - for path in image_paths: - path = path.resolve() - _ensure_regular_file(path, label="Qwen2.5 PI input image") - try: - with Image.open(path) as opened: - opened.load() - image = opened.convert("RGB") - except (OSError, ValueError) as exc: - raise StarVLAError(f"failed to decode input image {path}: {exc}") from exc - if image.size != (224, 224): - raise StarVLAError( - "Qwen2.5 PI parity requires an already-224x224 image so the " - "released deployment pre-resize is unambiguous" - ) - images.append(image) - records.append( - { - "source_path": str(path), - "source_size": path.stat().st_size, - "source_sha256": sha256_file(path), - "decoded_mode": image.mode, - "decoded_size": list(image.size), - "decoded_pixel_sha256": _image_pixel_sha256(image), - } - ) - if len(images) != 1: - raise StarVLAError("official Qwen2.5 PI oracle requires exactly one image") - return images, records - - -def _render_model_prompt(framework: Any, image: Any, task: str) -> str: - instruction = EXPECTED_COT_TEMPLATE.replace("{instruction}", task) - messages = [{ - "role": "user", - "content": [ - {"type": "image", "image": image}, - {"type": "text", "text": instruction}, - ], - }] - rendered = framework.qwen_vl_interface.processor.apply_chat_template( - messages, tokenize=False, add_generation_prompt=True - ) - if not isinstance(rendered, str): - raise StarVLAError("Qwen2.5 processor returned a non-string prompt") - return rendered - - -def run_official_forward( - framework: Any, - *, - images: Sequence[Any], - task: str, - state: np.ndarray, -) -> dict[str, Any]: - """Run Qwen and the exact legacy 16-block action head with explicit noise.""" - - import torch - - if len(images) != 1: - raise StarVLAError("legacy PI forward requires one image") - if state.shape != (1, 1, EXPECTED_STATE_DIM): - raise StarVLAError(f"legacy PI state must have shape [1,1,7], got {state.shape}") - qwen = framework.qwen_vl_interface - action_model = framework.action_model - qwen_inputs = qwen.build_qwenvl_inputs( - images=[list(images)], instructions=[task] - ) - with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16): - outputs = qwen( - **qwen_inputs, - output_attentions=False, - output_hidden_states=True, - return_dict=True, - ) - hidden = outputs.hidden_states - if hidden is None or len(hidden) != EXPECTED_QWEN_LAYER_COUNT + 1: - raise StarVLAError( - "Qwen2.5 output must expose hidden tuple indices 0..36" - ) - selected = list(hidden[-EXPECTED_DIT_BLOCK_COUNT:]) - if len(selected) != EXPECTED_DIT_BLOCK_COUNT: - raise StarVLAError("legacy PI did not select 16 conditioning states") - if selected[-1].dtype != torch.bfloat16: - raise StarVLAError("Qwen2.5 result_norm boundary must be BF16") - - state_tensor = torch.from_numpy(state).to( - device=selected[-1].device, dtype=torch.bfloat16 - ) - noise = explicit_initial_noise() - noise_tensor = torch.from_numpy(noise).to( - device=selected[-1].device, dtype=torch.bfloat16 - ) - original_randn = torch.randn - noise_calls = 0 - - def explicit_randn(*args: Any, **kwargs: Any): - nonlocal noise_calls - requested_size = kwargs.get("size", args[0] if args else None) - if tuple(requested_size) != tuple(noise_tensor.shape): - raise StarVLAError( - f"legacy PI requested unexpected noise shape: {requested_size}" - ) - if kwargs.get("dtype") != torch.bfloat16: - raise StarVLAError( - f"legacy PI requested unexpected noise dtype: {kwargs.get('dtype')}" - ) - noise_calls += 1 - if noise_calls != 1: - raise StarVLAError("legacy PI requested initial noise more than once") - return noise_tensor.clone() - - torch.randn = explicit_randn - try: - with torch.inference_mode(): - normalized_tensor = action_model.predict_action( - selected, state_tensor - ) - finally: - torch.randn = original_randn - if noise_calls != 1: - raise StarVLAError("legacy PI did not consume the explicit initial noise") - normalized = np.ascontiguousarray( - normalized_tensor.detach().cpu().float().numpy(), dtype=np.float32 - ) - if ( - normalized.shape - != (1, EXPECTED_ACTION_HORIZON, EXPECTED_ACTION_DIM) - or not np.isfinite(normalized).all() - ): - raise StarVLAError( - f"legacy PI produced invalid normalized actions: {normalized.shape}" - ) - qwen_arrays = { - key: value.detach().cpu() - for key, value in qwen_inputs.items() - if isinstance(value, torch.Tensor) - } - required = {"input_ids", "attention_mask", "image_grid_thw"} - if not required.issubset(qwen_arrays): - raise StarVLAError( - f"Qwen2.5 processor inputs are missing: {sorted(required - qwen_arrays.keys())}" - ) - return { - "qwen_inputs": qwen_arrays, - "result_norm": selected[-1].detach(), - "selected_hidden_tuple_indices": EXPECTED_HIDDEN_TUPLE_INDICES, - "policy_conditioning_dtype": "bfloat16", - "initial_noise": noise, - "normalized_actions": normalized, - } - - -def _array_sha256(array: np.ndarray) -> str: - contiguous = np.ascontiguousarray(array) - header = _canonical_json( - {"dtype": contiguous.dtype.str, "shape": list(contiguous.shape)} - ) - return _sha256_bytes(header + b"\x00" + contiguous.tobytes(order="C")) - - -def write_golden( - *, - output_dir: Path, - paths: Mapping[str, Any], - framework: Any, - image_path: Path, - image_record: Mapping[str, Any], - image: Any, - task: str, - unnorm_key: str, - state: np.ndarray, - captures: Mapping[str, Any], - unnormalized: np.ndarray, -) -> Path: - import torch - import transformers - - output_dir = output_dir.resolve() - if output_dir.exists(): - raise StarVLAError(f"golden output directory already exists: {output_dir}") - output_dir.parent.mkdir(parents=True, exist_ok=True) - - noise = np.ascontiguousarray(captures["initial_noise"], dtype=" np.ndarray: - try: - values = [float(item.strip()) for item in value.split(",")] - except ValueError as exc: - raise argparse.ArgumentTypeError( - "--state must contain seven comma-separated finite floats" - ) from exc - if len(values) != EXPECTED_STATE_DIM or not all( - math.isfinite(item) for item in values - ): - raise argparse.ArgumentTypeError( - "--state must contain seven comma-separated finite floats" - ) - return np.asarray(values, dtype=np.float32).reshape(1, 1, EXPECTED_STATE_DIM) - - -def _preflight_record( - paths: Mapping[str, Any], processor: Mapping[str, Any] -) -> dict[str, Any]: - noise = explicit_initial_noise() - return { - "schema_version": SCHEMA_VERSION, - "kind": "starvla_qwen25_pi_preflight", - "variant": VARIANT, - "model_type": MODEL_TYPE, - "backbone": BACKBONE, - "checkpoint": str(paths["checkpoint"]), - "checkpoint_ready": paths["checkpoint_ready"], - "expected_checkpoint": { - "bundle_uuid": official_bundle_uuid( - paths["variant"], paths["catalog"] - ), - "repo_id": OFFICIAL_CHECKPOINT_REPO_ID, - "revision": OFFICIAL_CHECKPOINT_REVISION, - "filename": OFFICIAL_CHECKPOINT_FILENAME, - "size": OFFICIAL_CHECKPOINT_SIZE, - "sha256": OFFICIAL_CHECKPOINT_SHA256, - }, - "qwen": {**paths["tokenizer"], "processor": dict(processor)}, - "conditioning": { - "hidden_tuple_indices": EXPECTED_HIDDEN_TUPLE_INDICES, - "terminal_tap": "result_norm", - "hidden_size": EXPECTED_QWEN_HIDDEN_DIM, - "transport": "native_bfloat16", - }, - "policy": { - "block_count": EXPECTED_DIT_BLOCK_COUNT, - "block_mode": "layerwise_cross_attention_every_block", - "released_config_interleave_self_attention": True, - "use_canonical_dit_forward": False, - "reference_execution_precision": - "whole_model_bf16_via_use_bf16", - "state_dim": EXPECTED_STATE_DIM, - "parameter_dtype": "bfloat16", - "runtime_contract_sha256": PI_RUNTIME_CONTRACT_SHA256, - }, - "action": { - "shape": [1, EXPECTED_ACTION_HORIZON, EXPECTED_ACTION_DIM], - "initial_noise_dtype": "bfloat16", - "initial_noise_algorithm": NOISE_ALGORITHM, - "initial_noise_array_sha256": _array_sha256(noise), - "timestep_ids": EXPECTED_TIMESTEP_IDS, - }, - "image": { - "required_parity_fixture_size": [224, 224], - "reason": - "avoid ambiguity in the released deployment pre-resize path", - }, - "action_gate": { - "reference": "local_original_checkpoint_python", - "metric": "full_tensor_global_relative_l2", - "operator": "<=", - "limit": ACTION_RELATIVE_L2_LIMIT, - "required_outputs": ["normalized_actions", "unnormalized_actions"], - }, - "source_probe": paths["source_probe"], - "effective_config_valid": True, - "golden_created": False, - } - - -def build_parser() -> argparse.ArgumentParser: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--checkpoint-root", type=Path, default=Path("ckpts/starvla")) - parser.add_argument("--checkpoint", type=Path) - parser.add_argument("--qwen-model", type=Path) - parser.add_argument( - "--starvla-source", - type=Path, - default=Path("ckpts/starvla/source/starvla"), - ) - parser.add_argument("--catalog", type=Path, default=DEFAULT_CATALOG) - parser.add_argument("--image", action="append", default=[], type=Path) - parser.add_argument( - "--task", default="put small spoon from basket to tray" - ) - parser.add_argument("--unnorm-key", choices=UNNORM_KEYS, default="oxe_bridge") - parser.add_argument( - "--state", - type=_parse_state, - default=_parse_state("0,0,0,0,0,0,0"), - help="seven comma-separated model-space state values", - ) - parser.add_argument("--device", default="cuda:0") - parser.add_argument( - "--output-dir", - type=Path, - default=Path( - "goldens/starvla/qwen25_pi/" - "bridge-episode-000000-frame000-put-spoon" - ), - ) - parser.add_argument("--preflight", "--preflight-only", action="store_true") - return parser - - -def main(argv: Sequence[str] | None = None) -> int: - args = build_parser().parse_args(argv) - holder: tempfile.TemporaryDirectory[str] | None = None - try: - _require_isolated_python() - import torch - import transformers - - validate_runtime_versions( - torch_version=torch.__version__, - torchvision_version=_distribution_version("torchvision"), - transformers_version=transformers.__version__, - numpy_version=np.__version__, - ) - if _distribution_version("qwen-vl-utils") != EXPECTED_QWEN_VL_UTILS_VERSION: - raise StarVLAError( - "qwen-vl-utils must be " - f"{EXPECTED_QWEN_VL_UTILS_VERSION} for the official oracle" - ) - _configure_determinism(torch, seed=0, device=args.device) - paths = validate_local_inputs( - checkpoint_root=args.checkpoint_root, - checkpoint=args.checkpoint, - qwen_model=args.qwen_model, - source_dir=args.starvla_source, - catalog_path=args.catalog, - ) - processor = validate_processor_contract(Path(paths["qwen_dir"])) - if args.preflight: - print( - json.dumps( - _preflight_record(paths, processor), - allow_nan=False, - indent=2, - sort_keys=True, - ) - ) - return 0 - if not paths["checkpoint_ready"]: - raise StarVLAError( - f"official Qwen2.5 PI checkpoint is not ready: {paths['checkpoint']}" - ) - if len(args.image) != 1: - raise StarVLAError("exactly one --image is required") - images, image_records = _load_images(args.image) - framework, _config, holder = load_official_framework( - paths, device=args.device - ) - captures = run_official_forward( - framework, - images=images, - task=args.task, - state=args.state, - ) - unnormalized = unnormalize_actions( - captures["normalized_actions"], paths["norm_stats"], args.unnorm_key - ) - manifest = write_golden( - output_dir=args.output_dir, - paths=paths, - framework=framework, - image_path=args.image[0].resolve(), - image_record=image_records[0], - image=images[0], - task=args.task, - unnorm_key=args.unnorm_key, - state=args.state, - captures=captures, - unnormalized=unnormalized, - ) - print(f"Wrote StarVLA Qwen2.5 PI local-Python golden: {manifest}") - return 0 - except (StarVLAError, OSError, RuntimeError, ValueError) as exc: - print(f"error: {exc}", file=sys.stderr) - return 2 - finally: - if holder is not None: - holder.cleanup() - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/tools/hf2gguf/starvla/groot_golden_constraints.txt b/tools/hf2gguf/starvla/groot_golden_constraints.txt deleted file mode 100644 index 12016b3..0000000 --- a/tools/hf2gguf/starvla/groot_golden_constraints.txt +++ /dev/null @@ -1,13 +0,0 @@ -# Numeric runtime frozen by the official StarVLA checkpoint environment. -# Apply this as constraints on top of the pinned StarVLA requirements when -# generating the independent Qwen-GR00T fixed-noise oracle. -torch==2.6.0 -torchvision==0.21.0 -transformers==4.57.0 -numpy==1.26.4 -diffusers==0.37.1 -tokenizers==0.22.2 -pillow==12.1.1 -omegaconf==2.3.0 -accelerate==1.5.2 -safetensors==0.7.0 diff --git a/tools/hf2gguf/starvla/inspect_starvla_checkpoint.py b/tools/hf2gguf/starvla/inspect_starvla_checkpoint.py deleted file mode 100755 index 822ebf0..0000000 --- a/tools/hf2gguf/starvla/inspect_starvla_checkpoint.py +++ /dev/null @@ -1,177 +0,0 @@ -#!/usr/bin/env python3 -"""Inspect a StarVLA .pt checkpoint without materializing copied weights.""" - -from __future__ import annotations - -import argparse -import json -import sys -from pathlib import Path -from typing import Any - -from starvla_checkpoint import ( - DEFAULT_CATALOG, - StarVLAError, - atomic_write_json, - build_inventory, - get_variant, - inventory_summary, - load_catalog, - load_checkpoint_state, - resolve_effective_config, - sha256_file, - verify_catalog_files, -) - - -def _load_structured(path: Path) -> dict[str, Any]: - if path.suffix == ".json": - return json.loads(path.read_text(encoding="utf-8")) - try: - import yaml - except ImportError as exc: - raise StarVLAError("PyYAML is required to inspect StarVLA YAML configs") from exc - value = yaml.safe_load(path.read_text(encoding="utf-8")) - if not isinstance(value, dict): - raise StarVLAError(f"expected an object in config file {path}") - return value - - -def _flatten(value: Any, prefix: str = "") -> dict[str, Any]: - if not isinstance(value, dict): - return {prefix: value} - flattened: dict[str, Any] = {} - for key, child in value.items(): - path = f"{prefix}.{key}" if prefix else str(key) - flattened.update(_flatten(child, path)) - return flattened - - -def inspect_config_candidates(source_dir: Path | None) -> dict[str, Any]: - if source_dir is None: - return {"files": {}, "conflicts": {}} - configs: dict[str, dict[str, Any]] = {} - for name in ("config.json", "config.yaml", "config.full.yaml"): - path = source_dir / name - if path.is_file(): - configs[name] = _load_structured(path) - - flattened = {name: _flatten(value) for name, value in configs.items()} - all_keys = sorted({key for values in flattened.values() for key in values}) - conflicts: dict[str, dict[str, Any]] = {} - for key in all_keys: - observed = {name: values[key] for name, values in flattened.items() if key in values} - serialized = {json.dumps(value, sort_keys=True) for value in observed.values()} - if len(observed) > 1 and len(serialized) > 1: - conflicts[key] = observed - return {"files": configs, "conflicts": conflicts} - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("checkpoint", type=Path) - parser.add_argument( - "--variant", - required=True, - choices=( - "oft", - "groot", - "pi_v3", - "qwen25_oft", - "qwen25_groot", - "qwen25_pi", - ), - ) - parser.add_argument("--catalog", type=Path, default=DEFAULT_CATALOG) - parser.add_argument("--source-dir", type=Path, help="directory containing config/statistics files") - parser.add_argument("--output", type=Path, required=True) - parser.add_argument("--effective-config-output", type=Path) - parser.add_argument("--skip-hash-check", action="store_true") - parser.add_argument( - "--allow-nonofficial-inventory", - action="store_true", - help="do not enforce pinned tensor counts; intended only for synthetic tests", - ) - return parser.parse_args() - - -def main() -> int: - args = parse_args() - try: - effective_config_output = args.effective_config_output or args.output.with_name( - "effective_config.json" - ) - output_targets = [args.output] - if args.source_dir is not None: - output_targets.append(effective_config_output) - if len(set(output_targets)) != len(output_targets): - raise StarVLAError("inspection and effective-config outputs must be different files") - for output_target in output_targets: - if output_target.exists() or output_target.is_symlink(): - raise StarVLAError(f"refusing to overwrite existing output: {output_target}") - - catalog = load_catalog(args.catalog) - variant = get_variant(catalog, args.variant) - if not args.checkpoint.is_file(): - raise StarVLAError(f"checkpoint does not exist: {args.checkpoint}") - actual_checkpoint = { - "path": str(args.checkpoint.resolve()), - "size": args.checkpoint.stat().st_size, - "sha256": sha256_file(args.checkpoint), - } - expected_checkpoint = variant["checkpoint"] - checkpoint_verified = ( - actual_checkpoint["size"] == expected_checkpoint["size"] - and actual_checkpoint["sha256"] == expected_checkpoint["sha256"] - ) - if not args.skip_hash_check and not checkpoint_verified: - raise StarVLAError( - f"checkpoint size/SHA256 mismatch for {args.checkpoint}: " - f"expected {expected_checkpoint['size']}/{expected_checkpoint['sha256']}, " - f"got {actual_checkpoint['size']}/{actual_checkpoint['sha256']}" - ) - source_assets_verified = False - if args.source_dir is not None: - verify_catalog_files(args.source_dir, variant) - source_assets_verified = True - state_dict = load_checkpoint_state(args.checkpoint) - records = build_inventory( - state_dict, - variant, - enforce_expected=not args.allow_nonofficial_inventory, - ) - effective_config = ( - resolve_effective_config(args.source_dir, args.variant, variant) - if args.source_dir - else None - ) - if effective_config is not None: - atomic_write_json(effective_config_output, effective_config, overwrite=False) - result = { - "schema_version": 1, - "variant": args.variant, - "model_type": variant["model_type"], - "source": { - "repo_id": variant["repo_id"], - "revision": variant["revision"], - "catalog_checkpoint": expected_checkpoint, - "input_checkpoint": actual_checkpoint, - "checkpoint_verification": "verified" if checkpoint_verified else "skipped_nonofficial", - "source_assets_verified": source_assets_verified, - }, - "summary": inventory_summary(records), - "config_candidates": inspect_config_candidates(args.source_dir), - "effective_config": str(effective_config_output) if effective_config is not None else None, - "tensors": [record.to_json() for record in records], - } - atomic_write_json(args.output, result, overwrite=False) - print(json.dumps(result["summary"], indent=2, sort_keys=True)) - print(f"inventory: {args.output}") - return 0 - except (StarVLAError, OSError, json.JSONDecodeError) as exc: - print(f"error: {exc}", file=sys.stderr) - return 1 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/tools/hf2gguf/starvla/pi_v3_golden_constraints.txt b/tools/hf2gguf/starvla/pi_v3_golden_constraints.txt deleted file mode 100644 index d0011fb..0000000 --- a/tools/hf2gguf/starvla/pi_v3_golden_constraints.txt +++ /dev/null @@ -1,14 +0,0 @@ -# Numeric runtime frozen by the official PI_v3 checkpoint's W&B run. -# Apply this as constraints on top of the pinned StarVLA requirements. -# HF path: wandb/wandb/run-20260426_011111-enstjn5q/files/requirements.txt -# SHA256: de6b505238663ea8a218620e8a4f99cbcfe1e6e09f347ab26f68fe434f3fb00e -torch==2.6.0 -torchvision==0.21.0 -transformers==4.57.0 -numpy==1.26.4 -diffusers==0.37.1 -tokenizers==0.22.2 -pillow==12.1.1 -omegaconf==2.3.0 -accelerate==1.5.2 -safetensors==0.7.0 diff --git a/tools/hf2gguf/starvla/serve_starvla_groot_reference.py b/tools/hf2gguf/starvla/serve_starvla_groot_reference.py deleted file mode 100644 index d17ed3e..0000000 --- a/tools/hf2gguf/starvla/serve_starvla_groot_reference.py +++ /dev/null @@ -1,535 +0,0 @@ -#!/usr/bin/env python3 -"""Serve the pinned official StarVLA Qwen3-VL GR00T Bridge checkpoint.""" - -from __future__ import annotations - -import argparse -import json -import logging -import sys -import time -from pathlib import Path -from typing import Any, Callable, Mapping, Sequence - -import numpy as np - - -TOOLS_DIR = Path(__file__).resolve().parent -REPO_ROOT = TOOLS_DIR.parents[2] -for search_path in (TOOLS_DIR, REPO_ROOT): - if str(search_path) not in sys.path: - sys.path.insert(0, str(search_path)) - -from generate_starvla_groot_golden import ( # noqa: E402 - EXPECTED_ACCELERATE_VERSION, - EXPECTED_DIFFUSERS_VERSION, - EXPECTED_NUMPY_VERSION, - EXPECTED_OMEGACONF_VERSION, - EXPECTED_PILLOW_VERSION, - EXPECTED_SAFETENSORS_VERSION, - EXPECTED_TOKENIZERS_VERSION, - EXPECTED_TORCHVISION_VERSION, - EXPECTED_TORCH_VERSION, - EXPECTED_TRANSFORMERS_VERSION, - _assert_module_origin, - _configure_determinism, - _distribution_version, - load_official_framework, - validate_available_inputs, - validate_runtime_versions, -) # noqa: E402 -from serve_starvla_oft_reference import ( # noqa: E402 - DEFAULT_IMAGE_NAME, - REFERENCE_BACKEND, - REFERENCE_PURPOSE, - SERVER_METADATA_SCHEMA_VERSION, - PredictRequest, - PredictResult, - ProtocolError, - ReferenceProtocolServer, - _asset_manifest, - _canonical_json_bytes, - _canonical_sha256, - _git_tracked_index_sha256, - _git_tree_sha1, - build_runtime_metadata, - wire, - write_metadata, -) -from starvla_checkpoint import ( # noqa: E402 - DEFAULT_CATALOG, - StarVLAError, - official_bundle_uuid, - sha256_file, -) - - -MODEL_TYPE = "starvla" -FRAMEWORK = "groot" -SUPPORTED_VARIANT = "groot" -DEFAULT_UNNORM_KEY = "oxe_bridge" -EXPECTED_PROFILES = ["oxe_bridge", "oxe_rt1"] - - -def _require_isolated_python() -> None: - if not sys.flags.isolated: - raise StarVLAError( - "the GR00T reference server must run in isolated mode; invoke with `python -I`" - ) - - -def _validate_reference_runtime(torch: Any, transformers: Any) -> None: - validate_runtime_versions( - torch_version=torch.__version__, - torchvision_version=_distribution_version("torchvision"), - transformers_version=transformers.__version__, - numpy_version=np.__version__, - diffusers_version=_distribution_version("diffusers"), - tokenizers_version=_distribution_version("tokenizers"), - pillow_version=_distribution_version("Pillow"), - omegaconf_version=_distribution_version("omegaconf"), - accelerate_version=_distribution_version("accelerate"), - safetensors_version=_distribution_version("safetensors"), - ) - - -def install_groot_dtype_bridge(framework: Any, torch: Any) -> None: - """Install the two explicit widens used by the independent GR00T oracle.""" - - if getattr(framework, "_robotcpp_groot_dtype_bridge", False): - return - action_model = framework.action_model - policy_dtype = next(action_model.parameters()).dtype - if policy_dtype != torch.float32: - raise StarVLAError(f"official GR00T policy must remain float32, got {policy_dtype}") - - original_action_encoder = action_model.action_encoder.forward - original_dit = action_model.model.forward - - def action_encoder_with_policy_dtype(actions: Any, timesteps: Any): - if actions.dtype not in (torch.bfloat16, torch.float32): - raise StarVLAError(f"unexpected GR00T action dtype: {actions.dtype}") - return original_action_encoder(actions.to(dtype=torch.float32), timesteps) - - def dit_with_policy_dtype(*args: Any, **kwargs: Any): - conditioning = kwargs.get( - "encoder_hidden_states", args[1] if len(args) > 1 else None - ) - if conditioning is None: - raise StarVLAError("official GR00T DiT omitted encoder_hidden_states") - if conditioning.dtype != torch.bfloat16: - raise StarVLAError( - f"official GR00T Qwen conditioning must be bfloat16, got {conditioning.dtype}" - ) - if "encoder_hidden_states" in kwargs: - kwargs["encoder_hidden_states"] = conditioning.to(dtype=torch.float32) - else: - mutable_args = list(args) - mutable_args[1] = conditioning.to(dtype=torch.float32) - args = tuple(mutable_args) - return original_dit(*args, **kwargs) - - action_model.action_encoder.forward = action_encoder_with_policy_dtype - action_model.model.forward = dit_with_policy_dtype - framework._robotcpp_groot_dtype_bridge = True - - -def build_preflight_record(paths: Mapping[str, Any]) -> dict[str, Any]: - variant = paths["variant"] - qwen = paths["qwen"] - return { - "schema_version": 1, - "ready": bool(paths["checkpoint_ready"]), - "variant": variant["_catalog_key"], - "model_type": variant["model_type"], - "framework": variant["framework"], - "backbone": variant.get("backbone", "qwen3_vl"), - "checkpoint": { - "repo_id": variant["repo_id"], - "revision": variant["revision"], - "path": str(Path(paths["checkpoint"]).resolve()), - "size": int(variant["checkpoint"]["size"]), - "sha256": variant["checkpoint"]["sha256"], - }, - "qwen": { - "repo_id": qwen["repo_id"], - "revision": qwen["revision"], - "path": str(Path(paths["qwen_dir"]).resolve()), - }, - "starvla": { - "revision": paths["catalog"]["source_revisions"]["starvla"], - "path": str(Path(paths["source_dir"]).resolve()), - }, - "catalog": { - "path": str(Path(paths["catalog_path"]).resolve()), - "sha256": sha256_file(Path(paths["catalog_path"])), - }, - } - - -def build_server_metadata( - paths: Mapping[str, Any], - framework: Any, - *, - default_unnorm_key: str, - source_tree_sha1: str, - source_tracked_index_sha256: str, - runtime: Mapping[str, Any], -) -> dict[str, Any]: - catalog = paths["catalog"] - variant = paths["variant"] - qwen = paths["qwen"] - statistics_path = Path(paths["policy_dir"]) / "dataset_statistics.json" - try: - statistics = json.loads(statistics_path.read_text(encoding="utf-8")) - except (OSError, json.JSONDecodeError) as exc: - raise StarVLAError( - f"failed to read official GR00T normalization profiles: {exc}" - ) from exc - if not isinstance(statistics, Mapping): - raise StarVLAError("official GR00T dataset_statistics.json must be an object") - profiles = [str(value) for value in statistics.keys()] - if profiles != EXPECTED_PROFILES: - raise StarVLAError( - f"unexpected GR00T normalization profiles: expected {EXPECTED_PROFILES}, got {profiles}" - ) - if default_unnorm_key not in profiles: - raise StarVLAError( - f"default unnorm key {default_unnorm_key!r} is not in {profiles}" - ) - - qwen_dtypes = sorted( - { - str(parameter.dtype).removeprefix("torch.") - for parameter in framework.qwen_vl_interface.parameters() - } - ) - policy_dtypes = sorted( - { - str(parameter.dtype).removeprefix("torch.") - for parameter in framework.action_model.parameters() - } - ) - if qwen_dtypes != ["bfloat16"] or policy_dtypes != ["float32"]: - raise StarVLAError( - f"unexpected loaded dtype profile: qwen={qwen_dtypes}, groot={policy_dtypes}" - ) - - chunk_size = int(framework.action_horizon) - action_dim = int(framework.action_model.action_dim) - if (chunk_size, action_dim) != (16, 7): - raise StarVLAError( - f"unexpected official GR00T action contract: chunk={chunk_size}, dim={action_dim}" - ) - - checkpoint_sha256 = str(variant["checkpoint"]["sha256"]) - checkpoint_revision = str(variant["revision"]) - starvla_revision = str(catalog["source_revisions"]["starvla"]) - qwen_manifest = _asset_manifest(qwen) - policy_manifest = _asset_manifest(variant) - model_info = { - "model_type": MODEL_TYPE, - "framework": FRAMEWORK, - "bundle_uuid": official_bundle_uuid(variant, catalog), - "checkpoint_sha256": checkpoint_sha256, - "checkpoint_revision": checkpoint_revision, - "qwen_revision": str(qwen["revision"]), - "starvla_revision": starvla_revision, - "image_names": [DEFAULT_IMAGE_NAME], - "state_supported": False, - "state_dimension_dynamic": False, - "state_dim": 0, - "chunk_size": chunk_size, - "action_dim": action_dim, - "normalization_profiles": profiles, - "default_unnorm_key": default_unnorm_key, - } - return { - "schema_version": SERVER_METADATA_SCHEMA_VERSION, - "protocol_version": wire.VERSION, - "backend": REFERENCE_BACKEND, - "purpose": REFERENCE_PURPOSE, - "catalog_variant": SUPPORTED_VARIANT, - "backbone": "qwen3_vl", - "runtime": dict(runtime), - "model_info": model_info, - "checkpoint": { - "repo_id": variant["repo_id"], - "revision": checkpoint_revision, - "path": str(Path(paths["checkpoint"]).resolve()), - "size": int(variant["checkpoint"]["size"]), - "sha256": checkpoint_sha256, - "asset_manifest_sha256": _canonical_sha256(policy_manifest), - }, - "qwen": { - "repo_id": qwen["repo_id"], - "revision": qwen["revision"], - "bootstrap_assets_manifest_sha256": _canonical_sha256(qwen_manifest), - "bootstrap_assets": qwen_manifest["files"], - }, - "starvla_source": { - "revision": starvla_revision, - "commit_sha": starvla_revision, - "git_tree_sha1": source_tree_sha1, - "tracked_index_manifest_sha256": source_tracked_index_sha256, - "path": str(Path(paths["source_dir"]).resolve()), - }, - "catalog": { - "path": str(Path(paths["catalog_path"]).resolve()), - "sha256": sha256_file(Path(paths["catalog_path"])), - }, - "dtype_profile": { - "qwen_parameters": "bfloat16", - "qwen_conditioning": "bfloat16", - "action_noise_initial": "bfloat16", - "action_encoder_input_cast": "float32", - "dit_conditioning_cast": "float32", - "groot_parameters": "float32", - "wire_actions": "float32", - "whole_model_cast": False, - }, - "action_contract": {"chunk_size": chunk_size, "action_dim": action_dim}, - "normalization": { - "implementation": "official PolicyNormProcessor", - "available_unnorm_keys": profiles, - "default_unnorm_key": default_unnorm_key, - "runtime_robot_profile_aliases": {profile: profile for profile in profiles}, - }, - "runtime_version_contract": { - "torch": EXPECTED_TORCH_VERSION, - "torchvision": EXPECTED_TORCHVISION_VERSION, - "transformers": EXPECTED_TRANSFORMERS_VERSION, - "numpy": EXPECTED_NUMPY_VERSION, - "diffusers": EXPECTED_DIFFUSERS_VERSION, - "tokenizers": EXPECTED_TOKENIZERS_VERSION, - "pillow": EXPECTED_PILLOW_VERSION, - "omegaconf": EXPECTED_OMEGACONF_VERSION, - "accelerate": EXPECTED_ACCELERATE_VERSION, - "safetensors": EXPECTED_SAFETENSORS_VERSION, - }, - } - - -class PinnedGROOTReferencePolicy: - """Original-checkpoint GR00T inference plus the official unnormalizer.""" - - def __init__( - self, - *, - framework: Any, - processor_factory: Callable[..., Any], - checkpoint: Path, - metadata: Mapping[str, Any], - ) -> None: - self.framework = framework - self.metadata = dict(metadata) - self.model_info = dict(self.metadata["model_info"]) - self.unnorm_key = str(self.model_info["default_unnorm_key"]) - self.processor = processor_factory( - str(checkpoint), unnorm_key=self.unnorm_key - ) - if self.processor.unnorm_key != self.unnorm_key: - raise StarVLAError("PolicyNormProcessor selected the wrong profile") - - def reset(self) -> None: - # GR00T has no observation history. Preserve the seeded noise stream. - return None - - def predict(self, request: PredictRequest) -> PredictResult: - if len(request.images) != 1: - raise ProtocolError( - f"GR00T Bridge reference requires exactly one image, got {len(request.images)}" - ) - image = request.images[0] - if image.name != DEFAULT_IMAGE_NAME: - raise ProtocolError( - f"GR00T Bridge reference requires image name {DEFAULT_IMAGE_NAME!r}, got {image.name!r}" - ) - if request.state: - raise ProtocolError("GR00T Bridge reference does not accept robot state") - if not request.task.strip(): - raise ProtocolError("task must not be empty") - - try: - from PIL import Image - except ImportError as exc: - raise RuntimeError("Pillow is required for GR00T reference inference") from exc - - pil_image = Image.fromarray(image.to_rgb_array(), mode="RGB") - total_started = time.perf_counter() - forward_started = time.perf_counter() - output = self.framework.predict_action( - examples=[{"image": [pil_image], "lang": request.task}] - ) - forward_ms = (time.perf_counter() - forward_started) * 1000.0 - if not isinstance(output, Mapping) or "normalized_actions" not in output: - raise RuntimeError("official GR00T forward did not return normalized_actions") - normalized = np.asarray(output["normalized_actions"]) - expected_shape = ( - 1, - int(self.model_info["chunk_size"]), - int(self.model_info["action_dim"]), - ) - if normalized.shape != expected_shape or not np.isfinite(normalized).all(): - raise RuntimeError( - f"official GR00T returned invalid normalized actions: {normalized.shape}" - ) - - unnorm_started = time.perf_counter() - actions = np.asarray( - self.processor.unapply_actions(normalized[0]), - dtype=np.float32, - ) - unnorm_ms = (time.perf_counter() - unnorm_started) * 1000.0 - if actions.shape != expected_shape[1:] or not np.isfinite(actions).all(): - raise RuntimeError( - f"official PolicyNormProcessor returned invalid actions: {actions.shape}" - ) - return PredictResult( - actions=np.ascontiguousarray(actions), - metrics={ - "python_forward_ms": forward_ms, - "python_unnorm_ms": unnorm_ms, - "model_total_ms": (time.perf_counter() - total_started) * 1000.0, - }, - ) - - -def load_pinned_reference_policy( - *, - checkpoint_root: Path, - starvla_source: Path | None, - device: str, - noise_seed: int, - default_unnorm_key: str, -) -> PinnedGROOTReferencePolicy: - source_dir = starvla_source or checkpoint_root / "source" / "starvla" - paths = validate_available_inputs( - checkpoint_root=checkpoint_root, - source_dir=Path(source_dir), - catalog_path=DEFAULT_CATALOG, - ) - try: - import torch - import transformers - except ImportError as exc: - raise StarVLAError(f"official StarVLA runtime dependency is missing: {exc}") from exc - _validate_reference_runtime(torch, transformers) - _configure_determinism(torch, seed=noise_seed, device=device) - framework, _config = load_official_framework(paths, device=device) - install_groot_dtype_bridge(framework, torch) - - source_dir = Path(paths["source_dir"]) - sys.path.insert(0, str(source_dir)) - try: - from deployment.model_server import policy_norm_processor - - _assert_module_origin(policy_norm_processor, source_dir) - processor_factory = policy_norm_processor.PolicyNormProcessor - finally: - if sys.path and sys.path[0] == str(source_dir): - del sys.path[0] - - metadata = build_server_metadata( - paths, - framework, - default_unnorm_key=default_unnorm_key, - source_tree_sha1=_git_tree_sha1(source_dir), - source_tracked_index_sha256=_git_tracked_index_sha256(source_dir), - runtime=build_runtime_metadata(torch, transformers, device=device), - ) - return PinnedGROOTReferencePolicy( - framework=framework, - processor_factory=processor_factory, - checkpoint=Path(paths["checkpoint"]), - metadata=metadata, - ) - - -def build_parser() -> argparse.ArgumentParser: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--variant", choices=(SUPPORTED_VARIANT,), default=SUPPORTED_VARIANT) - parser.add_argument("--checkpoint-root", type=Path, default=Path("ckpts/starvla")) - parser.add_argument("--starvla-source", type=Path) - parser.add_argument("--device", default="cuda:0") - parser.add_argument("--host", default="127.0.0.1") - parser.add_argument("--port", type=int, default=5555) - parser.add_argument("--unnorm-key", default=DEFAULT_UNNORM_KEY) - parser.add_argument("--noise-seed", type=int, default=0) - parser.add_argument("--metadata-output", type=Path) - parser.add_argument("--preflight", action="store_true") - parser.add_argument("--verbosity", type=int, default=0) - return parser - - -def main(argv: Sequence[str] | None = None) -> int: - args = build_parser().parse_args(argv) - _require_isolated_python() - if args.host != "127.0.0.1": - raise StarVLAError("--host must be 127.0.0.1") - if args.port <= 0 or args.port > 65535: - raise StarVLAError("--port must be in 1..65535") - if args.noise_seed < 0: - raise StarVLAError("--noise-seed must be non-negative") - if args.verbosity < 0: - raise StarVLAError("--verbosity must be non-negative") - logging.basicConfig( - level=logging.DEBUG if args.verbosity else logging.INFO, - format="%(asctime)s %(levelname)s %(message)s", - force=True, - ) - - checkpoint_root = args.checkpoint_root.resolve() - source_dir = ( - args.starvla_source.resolve() - if args.starvla_source - else checkpoint_root / "source" / "starvla" - ) - if args.preflight: - paths = validate_available_inputs( - checkpoint_root=checkpoint_root, - source_dir=source_dir, - catalog_path=DEFAULT_CATALOG, - ) - record = build_preflight_record(paths) - if not record["ready"]: - raise StarVLAError(f"official GR00T checkpoint is incomplete: {paths['checkpoint']}") - if args.metadata_output is not None: - write_metadata(args.metadata_output, record) - sys.stdout.write(json.dumps(record, indent=2, sort_keys=True) + "\n") - return 0 - - policy = load_pinned_reference_policy( - checkpoint_root=checkpoint_root, - starvla_source=source_dir, - device=args.device, - noise_seed=args.noise_seed, - default_unnorm_key=args.unnorm_key, - ) - if args.metadata_output is not None: - write_metadata(args.metadata_output, policy.metadata) - logging.info( - "loaded pinned GR00T Python reference metadata=%s", - _canonical_json_bytes(policy.metadata).decode("ascii"), - ) - server = ReferenceProtocolServer(policy, host=args.host, port=args.port) - logging.info( - "Python reference server listening on %s:%d model=%s variant=%s", - server.address[0], - server.address[1], - MODEL_TYPE, - args.variant, - ) - try: - server.serve_forever() - except KeyboardInterrupt: - logging.info("Python reference server interrupted") - server.close() - return 0 - - -if __name__ == "__main__": - try: - raise SystemExit(main()) - except StarVLAError as exc: - raise SystemExit(f"error: {exc}") from exc diff --git a/tools/hf2gguf/starvla/serve_starvla_oft_reference.py b/tools/hf2gguf/starvla/serve_starvla_oft_reference.py deleted file mode 100644 index 99f0787..0000000 --- a/tools/hf2gguf/starvla/serve_starvla_oft_reference.py +++ /dev/null @@ -1,1152 +0,0 @@ -#!/usr/bin/env python3 -"""Serve a pinned official StarVLA OFT checkpoint over robot.cpp protocol v3. - -This process is the Python-reference backend for closed-loop parity evaluation. -It deliberately shares the robot.cpp client protocol and SimplerEnv adapter, so -the only policy variable is the original PyTorch checkpoint versus GGUF. -""" - -from __future__ import annotations - -import argparse -import hashlib -import json -import logging -import math -import os -import socket -import struct -import subprocess -import sys -import tempfile -import time -from dataclasses import dataclass -from pathlib import Path -from typing import Any, Callable, Mapping, Sequence - -import numpy as np - - -TOOLS_DIR = Path(__file__).resolve().parent -REPO_ROOT = TOOLS_DIR.parents[2] -for search_path in (TOOLS_DIR, REPO_ROOT): - if str(search_path) not in sys.path: - sys.path.insert(0, str(search_path)) - -from generate_starvla_oft_golden import ( # noqa: E402 - _assert_module_origin, - _configure_determinism, - _distribution_version, - _require_isolated_python, - load_official_framework as load_qwen3_official_framework, - validate_official_inputs as validate_qwen3_official_inputs, - validate_runtime_versions, -) -from generate_starvla_qwen25_oft_golden import ( # noqa: E402 - EXPECTED_QWEN_VL_UTILS_VERSION, - LEGACY_UNNORM_PROFILES, - legacy_normalization_contract, - load_official_framework as load_qwen25_official_framework, - unnormalize_legacy_actions, - validate_local_inputs as validate_qwen25_local_inputs, -) -from robot_client.python import model_client as wire # noqa: E402 -from starvla_checkpoint import ( # noqa: E402 - DEFAULT_CATALOG, - StarVLAError, - get_qwen_asset, - get_variant, - load_catalog, - official_bundle_uuid, - sha256_file, - verify_catalog_files, -) - - -SERVER_METADATA_SCHEMA_VERSION = 2 -MODEL_TYPE = "starvla" -FRAMEWORK = "oft" -REFERENCE_BACKEND = "local-python-checkpoint-reference" -REFERENCE_PURPOSE = "bridge-only" -DEFAULT_IMAGE_NAME = "image_0" -DEFAULT_UNNORM_KEY = "oxe_bridge" -QWEN25_DEFAULT_UNNORM_KEY = "bridge_dataset" -SUPPORTED_VARIANTS = ("oft", "qwen25_oft") - -STATUS_BAD_REQUEST = 1 -STATUS_BAD_VERSION = 3 -STATUS_PAYLOAD_TOO_BIG = 4 -STATUS_INTERNAL_ERROR = 5 -MAX_PAYLOAD_BYTES = 256 * 1024 * 1024 - - -class ProtocolError(ValueError): - """A malformed or contract-invalid protocol request.""" - - -class PayloadTooBig(ProtocolError): - pass - - -@dataclass(frozen=True) -class RequestHeader: - magic: int - version: int - header_size: int - op: int - flags: int - request_id: int - status: int - payload_len: int - reserved: int - - -@dataclass(frozen=True) -class WireImage: - name: str - width: int - height: int - channels: int - stride_bytes: int - data: bytes - - def to_rgb_array(self) -> np.ndarray: - rows = np.frombuffer( - self.data, - dtype=np.uint8, - count=self.stride_bytes * self.height, - ).reshape(self.height, self.stride_bytes) - packed = rows[:, : self.width * self.channels] - return np.ascontiguousarray(packed.reshape(self.height, self.width, self.channels)) - - -@dataclass(frozen=True) -class PredictRequest: - images: tuple[WireImage, ...] - state: tuple[float, ...] - task: str - - -@dataclass(frozen=True) -class PredictResult: - actions: np.ndarray - metrics: Mapping[str, float] - - -def _canonical_json_bytes(value: Any) -> bytes: - return json.dumps( - value, - ensure_ascii=True, - sort_keys=True, - separators=(",", ":"), - ).encode("utf-8") - - -def _canonical_sha256(value: Any) -> str: - return hashlib.sha256(_canonical_json_bytes(value)).hexdigest() - - -def _decode_utf8(value: bytes, label: str) -> str: - try: - return value.decode("utf-8", errors="strict") - except UnicodeDecodeError as exc: - raise ProtocolError(f"{label} is not valid UTF-8") from exc - - -def _asset_manifest(entry: Mapping[str, Any]) -> dict[str, Any]: - return { - "repo_id": entry["repo_id"], - "revision": entry["revision"], - "files": { - relative: { - "size": int(entry["file_hashes"][relative]["size"]), - "sha256": str(entry["file_hashes"][relative]["sha256"]), - } - for relative in sorted(entry.get("files", [])) - }, - } - - -def _git_tree_sha1(source_dir: Path) -> str: - try: - result = subprocess.run( - ["git", "-C", str(source_dir), "rev-parse", "HEAD^{tree}"], - check=True, - capture_output=True, - text=True, - ) - except (OSError, subprocess.CalledProcessError) as exc: - raise StarVLAError(f"failed to resolve pinned StarVLA Git tree: {exc}") from exc - value = result.stdout.strip() - if len(value) != 40 or any(ch not in "0123456789abcdef" for ch in value): - raise StarVLAError(f"invalid pinned StarVLA Git tree SHA1: {value!r}") - return value - - -def _git_tracked_index_sha256(source_dir: Path) -> str: - """Hash the clean checkout's modes, paths, and Git blob identities.""" - - try: - result = subprocess.run( - ["git", "-C", str(source_dir), "ls-files", "-s", "-z"], - check=True, - capture_output=True, - ) - except (OSError, subprocess.CalledProcessError) as exc: - raise StarVLAError( - f"failed to hash pinned StarVLA tracked-file manifest: {exc}" - ) from exc - if not result.stdout: - raise StarVLAError("pinned StarVLA tracked-file manifest is empty") - return hashlib.sha256(result.stdout).hexdigest() - - -def default_unnorm_key_for_variant(variant: str) -> str: - if variant == "oft": - return DEFAULT_UNNORM_KEY - if variant == "qwen25_oft": - return QWEN25_DEFAULT_UNNORM_KEY - raise StarVLAError( - f"unsupported OFT reference variant {variant!r}; expected one of {SUPPORTED_VARIANTS}" - ) - - -def validate_reference_inputs( - *, - checkpoint_root: Path, - source_dir: Path, - variant_name: str, - catalog_path: Path = DEFAULT_CATALOG, -) -> dict[str, Any]: - """Resolve and verify one official OFT checkpoint without accepting aliases.""" - - if variant_name == "oft": - return validate_qwen3_official_inputs( - checkpoint_root=checkpoint_root, - source_dir=source_dir, - catalog_path=catalog_path, - ) - if variant_name != "qwen25_oft": - raise StarVLAError( - f"unsupported OFT reference variant {variant_name!r}; " - f"expected one of {SUPPORTED_VARIANTS}" - ) - - catalog = load_catalog(catalog_path) - variant = get_variant(catalog, variant_name) - qwen_asset_name, qwen = get_qwen_asset(catalog, variant) - if ( - variant.get("framework") != FRAMEWORK - or variant.get("model_type") != MODEL_TYPE - or variant.get("backbone") != "qwen2_5_vl" - or qwen_asset_name != "qwen2_5_vl_3b_instruct" - ): - raise StarVLAError("catalog Qwen2.5 OFT identity is incompatible") - - checkpoint_root = checkpoint_root.resolve() - expected_source = (checkpoint_root / "source" / "starvla").resolve() - source_dir = source_dir.resolve() - if source_dir != expected_source: - raise StarVLAError( - f"StarVLA source must be the canonical checkout {expected_source}, got {source_dir}" - ) - - policy_dir = ( - checkpoint_root / "sources" / variant["directory"] / variant["revision"] - ) - qwen_dir = checkpoint_root / "sources" / qwen["directory"] / qwen["revision"] - checkpoint = policy_dir / variant["checkpoint"]["path"] - verify_catalog_files(policy_dir, variant) - verify_catalog_files(qwen_dir, qwen) - local = validate_qwen25_local_inputs( - checkpoint=checkpoint, - qwen_model=qwen_dir, - source_dir=source_dir, - expected_checkpoint_sha256=str(variant["checkpoint"]["sha256"]), - expected_checkpoint_size=int(variant["checkpoint"]["size"]), - expected_source_revision=str(catalog["source_revisions"]["starvla"]), - ) - return { - **local, - "catalog": catalog, - "variant": variant, - "qwen": qwen, - "policy_dir": policy_dir, - "catalog_path": catalog_path.resolve(), - } - - -def build_preflight_record(paths: Mapping[str, Any]) -> dict[str, Any]: - variant = paths["variant"] - qwen = paths["qwen"] - return { - "schema_version": 1, - "ready": True, - "variant": variant["_catalog_key"], - "model_type": variant["model_type"], - "framework": variant["framework"], - "backbone": variant.get("backbone", "qwen3_vl"), - "checkpoint": { - "repo_id": variant["repo_id"], - "revision": variant["revision"], - "path": str(Path(paths["checkpoint"]).resolve()), - "size": int(variant["checkpoint"]["size"]), - "sha256": variant["checkpoint"]["sha256"], - }, - "qwen": { - "repo_id": qwen["repo_id"], - "revision": qwen["revision"], - "path": str(Path(paths["qwen_dir"]).resolve()), - }, - "starvla": { - "revision": paths["catalog"]["source_revisions"]["starvla"], - "path": str(Path(paths["source_dir"]).resolve()), - }, - "catalog": { - "path": str(Path(paths["catalog_path"]).resolve()), - "sha256": sha256_file(Path(paths["catalog_path"])), - }, - } - - -def apply_official_bf16(framework: Any, torch: Any) -> Any: - """Match the official Bridge server's ``--use_bf16`` whole-model cast.""" - - framework = framework.to(dtype=torch.bfloat16) - dtypes = {parameter.dtype for parameter in framework.parameters()} - if dtypes != {torch.bfloat16}: - raise StarVLAError(f"official OFT model must be entirely bfloat16, got {dtypes}") - return framework - - -def build_runtime_metadata( - torch: Any, - transformers: Any, - *, - device: str, -) -> dict[str, Any]: - """Record the exact local Python runtime used by the Bridge reference.""" - - torch_device = torch.device(device) - cudnn_version = torch.backends.cudnn.version() - record: dict[str, Any] = { - "python_full_version": sys.version, - "torch": str(torch.__version__), - "torch_cuda": ( - None if torch.version.cuda is None else str(torch.version.cuda) - ), - "cudnn": None if cudnn_version is None else int(cudnn_version), - "transformers": str(transformers.__version__), - "pillow": _distribution_version("Pillow"), - "numpy": str(np.__version__), - "device": str(torch_device), - "gpu_name": None, - "compute_capability": None, - } - if torch_device.type == "cuda": - index = ( - torch_device.index - if torch_device.index is not None - else torch.cuda.current_device() - ) - properties = torch.cuda.get_device_properties(index) - record["gpu_name"] = str(properties.name) - record["compute_capability"] = [ - int(properties.major), - int(properties.minor), - ] - return record - - -def build_server_metadata( - paths: Mapping[str, Any], - framework: Any, - *, - default_unnorm_key: str, - source_tree_sha1: str, - source_tracked_index_sha256: str, - runtime: Mapping[str, Any], -) -> dict[str, Any]: - catalog = paths["catalog"] - variant = paths["variant"] - qwen = paths["qwen"] - variant_name = str(variant["_catalog_key"]) - backbone = str(variant.get("backbone", "qwen3_vl")) - profiles = [str(value) for value in framework.norm_stats.keys()] - if not profiles or len(set(profiles)) != len(profiles): - raise StarVLAError(f"invalid official normalization profiles: {profiles}") - expected_profiles = ( - list(LEGACY_UNNORM_PROFILES) - if variant_name == "qwen25_oft" - else ["oxe_bridge", "oxe_rt1"] - ) - if profiles != expected_profiles: - raise StarVLAError( - f"unexpected {variant_name} normalization profiles: " - f"expected {expected_profiles}, got {profiles}" - ) - if default_unnorm_key not in profiles: - raise StarVLAError( - f"default unnorm key {default_unnorm_key!r} is not in {profiles}" - ) - - qwen_dtypes = sorted( - {str(parameter.dtype).removeprefix("torch.") for parameter in framework.qwen_vl_interface.parameters()} - ) - policy_dtypes = sorted( - {str(parameter.dtype).removeprefix("torch.") for parameter in framework.action_model.parameters()} - ) - if qwen_dtypes != ["bfloat16"] or policy_dtypes != ["bfloat16"]: - raise StarVLAError( - f"unexpected loaded dtype profile: qwen={qwen_dtypes}, oft={policy_dtypes}" - ) - chunk_size = int(framework.chunk_len) - action_dim = int(framework.action_model.action_dim) - if chunk_size <= 0 or action_dim <= 0: - raise StarVLAError( - f"invalid official action contract: chunk={chunk_size}, dim={action_dim}" - ) - - checkpoint_sha256 = str(variant["checkpoint"]["sha256"]) - checkpoint_revision = str(variant["revision"]) - qwen_manifest = _asset_manifest(qwen) - policy_manifest = _asset_manifest(variant) - starvla_revision = str(catalog["source_revisions"]["starvla"]) - model_info = { - "model_type": MODEL_TYPE, - "framework": FRAMEWORK, - "bundle_uuid": official_bundle_uuid(variant, catalog), - "checkpoint_sha256": checkpoint_sha256, - "checkpoint_revision": checkpoint_revision, - "qwen_revision": str(qwen["revision"]), - "starvla_revision": starvla_revision, - "image_names": [DEFAULT_IMAGE_NAME], - "state_supported": False, - "state_dimension_dynamic": False, - "state_dim": 0, - "chunk_size": chunk_size, - "action_dim": action_dim, - "normalization_profiles": profiles, - "default_unnorm_key": default_unnorm_key, - } - return { - "schema_version": SERVER_METADATA_SCHEMA_VERSION, - "protocol_version": wire.VERSION, - "backend": REFERENCE_BACKEND, - "purpose": REFERENCE_PURPOSE, - "catalog_variant": variant_name, - "backbone": backbone, - "runtime": dict(runtime), - "model_info": model_info, - "checkpoint": { - "repo_id": variant["repo_id"], - "revision": checkpoint_revision, - "path": str(Path(paths["checkpoint"]).resolve()), - "size": int(variant["checkpoint"]["size"]), - "sha256": checkpoint_sha256, - "asset_manifest_sha256": _canonical_sha256(policy_manifest), - }, - "qwen": { - "repo_id": qwen["repo_id"], - "revision": qwen["revision"], - "bootstrap_assets_manifest_sha256": _canonical_sha256(qwen_manifest), - "bootstrap_assets": qwen_manifest["files"], - }, - "starvla_source": { - "revision": starvla_revision, - "commit_sha": starvla_revision, - "git_tree_sha1": source_tree_sha1, - "tracked_index_manifest_sha256": source_tracked_index_sha256, - "path": str(Path(paths["source_dir"]).resolve()), - }, - "catalog": { - "path": str(Path(paths["catalog_path"]).resolve()), - "sha256": sha256_file(Path(paths["catalog_path"])), - }, - "dtype_profile": { - "qwen_parameters": "bfloat16", - "qwen_action_queries": "bfloat16", - "oft_input_cast": None, - "oft_parameters": "bfloat16", - "wire_actions": "float32", - "whole_model_cast": True, - }, - "action_contract": { - "chunk_size": chunk_size, - "action_dim": action_dim, - }, - "normalization": { - "implementation": ( - "released_q01_q99_masked_with_binary_unmasked_dimensions" - if variant_name == "qwen25_oft" - else "official PolicyNormProcessor" - ), - "available_unnorm_keys": profiles, - "default_unnorm_key": default_unnorm_key, - "runtime_robot_profile_aliases": ( - dict(LEGACY_UNNORM_PROFILES) - if variant_name == "qwen25_oft" - else {profile: profile for profile in profiles} - ), - }, - } - - -def decode_request_header(raw: bytes) -> RequestHeader: - if len(raw) != wire.HEADER_SIZE: - raise ProtocolError("short header") - return RequestHeader(*wire.HEADER.unpack(raw)) - - -def validate_request_header(header: RequestHeader) -> None: - if header.magic != wire.MAGIC: - raise ProtocolError("bad magic") - if header.version != wire.VERSION: - raise ProtocolError("bad protocol version") - if header.header_size != wire.HEADER_SIZE: - raise ProtocolError("bad header size") - if header.flags != 0 or header.status != wire.STATUS_OK or header.reserved != 0: - raise ProtocolError("request header flags/status/reserved must be zero") - if header.payload_len > MAX_PAYLOAD_BYTES: - raise PayloadTooBig("payload too large") - - -def decode_predict_request(payload: bytes) -> PredictRequest: - if len(payload) < wire.PREDICT_REQ_V2_FIXED.size: - raise ProtocolError("short predict request") - image_count, state_count, task_len = wire.PREDICT_REQ_V2_FIXED.unpack_from(payload) - if image_count == 0: - raise ProtocolError("predict request requires at least one image") - - offset = wire.PREDICT_REQ_V2_FIXED.size - remaining = len(payload) - offset - if image_count > remaining // wire.PREDICT_REQ_V2_IMAGE.size: - raise ProtocolError("image count exceeds predict request metadata") - - metadata: list[tuple[int, int, int, int, int, int]] = [] - for index in range(image_count): - ( - image_format, - name_len, - width, - height, - channels, - stride_bytes, - data_len, - ) = wire.PREDICT_REQ_V2_IMAGE.unpack_from(payload, offset) - offset += wire.PREDICT_REQ_V2_IMAGE.size - if image_format != wire.IMAGE_RAW_RGB_U8: - raise ProtocolError(f"image[{index}] has an unsupported image format") - if width <= 0 or height <= 0 or channels != 3: - raise ProtocolError(f"image[{index}] has invalid raw RGB dimensions") - packed_stride = width * channels - if stride_bytes < packed_stride: - raise ProtocolError(f"image[{index}] has an invalid stride_bytes") - if data_len < stride_bytes * height: - raise ProtocolError( - f"image[{index}] data is smaller than stride_bytes * height" - ) - metadata.append((name_len, width, height, channels, stride_bytes, data_len)) - - body_size = ( - state_count * 4 - + task_len - + sum(name_len + data_len for name_len, *_rest, data_len in metadata) - ) - if body_size != len(payload) - offset: - raise ProtocolError("predict request fields do not exactly match payload") - - state: tuple[float, ...] - if state_count: - state = tuple(struct.unpack_from(f"<{state_count}f", payload, offset)) - else: - state = () - offset += state_count * 4 - if any(not math.isfinite(value) for value in state): - raise ProtocolError("state contains a non-finite value") - - task = _decode_utf8(payload[offset : offset + task_len], "task") - offset += task_len - - images: list[WireImage] = [] - for index, (name_len, width, height, channels, stride_bytes, data_len) in enumerate(metadata): - name = _decode_utf8(payload[offset : offset + name_len], f"image[{index}] name") - offset += name_len - data = bytes(payload[offset : offset + data_len]) - offset += data_len - images.append( - WireImage( - name=name, - width=width, - height=height, - channels=channels, - stride_bytes=stride_bytes, - data=data, - ) - ) - if offset != len(payload): - raise ProtocolError("trailing bytes in predict request") - return PredictRequest( - images=tuple(images), - state=state, - task=task, - ) - - -def encode_predict_response(result: PredictResult) -> bytes: - actions = np.asarray(result.actions, dtype=np.float32) - if actions.ndim != 2 or actions.shape[0] <= 0 or actions.shape[1] <= 0: - raise ProtocolError(f"invalid action matrix: {actions.shape}") - if not np.isfinite(actions).all(): - raise ProtocolError("actions contain non-finite values") - chunk_size, action_dim = (int(actions.shape[0]), int(actions.shape[1])) - action_count = actions.size - metric_rows: list[tuple[bytes, float]] = [] - for name in sorted(result.metrics): - name_bytes = str(name).encode("utf-8") - if not name_bytes: - raise ProtocolError("metric name is empty") - value = float(result.metrics[name]) - if not math.isfinite(value): - raise ProtocolError(f"metric {name!r} is non-finite") - metric_rows.append((name_bytes, value)) - - output = bytearray( - wire.PREDICT_RESP_FIXED.pack( - chunk_size, - action_dim, - action_count, - len(metric_rows), - ) - ) - for name, value in metric_rows: - output += wire.PREDICT_RESP_METRIC.pack(len(name), value) - output += name - output += np.ascontiguousarray(actions, dtype=" None: - if (processor_factory is None) == (action_unnormalizer is None): - raise StarVLAError( - "reference policy requires exactly one action unnormalization implementation" - ) - self.framework = framework - self.action_unnormalizer = action_unnormalizer - self.metadata = dict(metadata) - self.model_info = dict(self.metadata["model_info"]) - self.unnorm_key = str(self.model_info["default_unnorm_key"]) - self.processor = None - if processor_factory is not None: - self.processor = processor_factory( - str(checkpoint), unnorm_key=self.unnorm_key - ) - if getattr(self.processor, "unnorm_key", self.unnorm_key) != self.unnorm_key: - raise StarVLAError("PolicyNormProcessor selected the wrong profile") - - def reset(self) -> None: - # OFT is stateless; this method intentionally preserves loaded weights. - return None - - def predict(self, request: PredictRequest) -> PredictResult: - if len(request.images) != 1: - raise ProtocolError( - f"OFT Bridge reference requires exactly one image, got {len(request.images)}" - ) - image = request.images[0] - if image.name != DEFAULT_IMAGE_NAME: - raise ProtocolError( - f"OFT Bridge reference requires image name {DEFAULT_IMAGE_NAME!r}, got {image.name!r}" - ) - if request.state: - raise ProtocolError("OFT Bridge reference does not accept robot state") - if not request.task.strip(): - raise ProtocolError("task must not be empty") - - try: - from PIL import Image - except ImportError as exc: - raise RuntimeError("Pillow is required for OFT reference inference") from exc - - pil_image = Image.fromarray(image.to_rgb_array(), mode="RGB") - total_started = time.perf_counter() - forward_started = time.perf_counter() - output = self.framework.predict_action( - examples=[{"image": [pil_image], "lang": request.task}] - ) - forward_ms = (time.perf_counter() - forward_started) * 1000.0 - if not isinstance(output, Mapping) or "normalized_actions" not in output: - raise RuntimeError("official OFT forward did not return normalized_actions") - normalized = np.asarray(output["normalized_actions"]) - expected_shape = ( - 1, - int(self.model_info["chunk_size"]), - int(self.model_info["action_dim"]), - ) - if normalized.shape != expected_shape or not np.isfinite(normalized).all(): - raise RuntimeError( - f"official OFT returned invalid normalized actions: {normalized.shape}" - ) - - unnorm_started = time.perf_counter() - if self.action_unnormalizer is None: - assert self.processor is not None - actions = np.asarray( - self.processor.unapply_actions(normalized[0]), - dtype=np.float32, - ) - else: - actions = np.asarray( - self.action_unnormalizer(normalized, self.unnorm_key), - dtype=np.float32, - ) - unnorm_ms = (time.perf_counter() - unnorm_started) * 1000.0 - if actions.shape != expected_shape[1:] or not np.isfinite(actions).all(): - raise RuntimeError( - f"official PolicyNormProcessor returned invalid actions: {actions.shape}" - ) - total_ms = (time.perf_counter() - total_started) * 1000.0 - return PredictResult( - actions=np.ascontiguousarray(actions), - metrics={ - "python_forward_ms": forward_ms, - "python_unnorm_ms": unnorm_ms, - "model_total_ms": total_ms, - }, - ) - - -class ProtocolApplication: - def __init__(self, policy: Any): - self.policy = policy - self.shutdown_requested = False - - def dispatch( - self, - op: int, - payload: bytes, - *, - server_recv_ms: float = 0.0, - ) -> tuple[bytes, bool]: - if op == wire.OP_HEALTH: - if payload: - raise ProtocolError("health request payload must be empty") - return f"ok policy={self.policy.model_info['model_type']}".encode(), False - if op == wire.OP_RESET: - if payload: - raise ProtocolError("reset request payload must be empty") - self.policy.reset() - return b"ok", False - if op == wire.OP_SHUTDOWN: - if payload: - raise ProtocolError("shutdown request payload must be empty") - self.shutdown_requested = True - return b"ok", True - if op == wire.OP_PREDICT: - request = decode_predict_request(payload) - predict_started = time.perf_counter() - result = self.policy.predict(request) - server_predict_ms = (time.perf_counter() - predict_started) * 1000.0 - metrics = dict(result.metrics) - metrics.update( - { - "server_queue_ms": 0.0, - "server_predict_ms": server_predict_ms, - "server_recv_ms": float(server_recv_ms), - } - ) - return ( - encode_predict_response( - PredictResult( - actions=result.actions, - metrics=metrics, - ) - ), - False, - ) - raise ProtocolError("unknown op") - - -def _recv_exact(sock: socket.socket, length: int, *, allow_initial_eof: bool = False) -> bytes | None: - chunks: list[bytes] = [] - remaining = length - while remaining: - chunk = sock.recv(remaining) - if not chunk: - if allow_initial_eof and remaining == length: - return None - raise ConnectionError("peer closed connection") - chunks.append(chunk) - remaining -= len(chunk) - return b"".join(chunks) - - -def _response_header( - request: RequestHeader, - *, - status: int, - payload_len: int, -) -> bytes: - return wire.HEADER.pack( - wire.MAGIC, - wire.VERSION, - wire.HEADER_SIZE, - request.op, - 0, - request.request_id, - status, - payload_len, - 0, - ) - - -class ReferenceProtocolServer: - """Small sequential TCP server matching robot_server/session.cpp semantics.""" - - def __init__( - self, - policy: Any, - *, - host: str = "127.0.0.1", - port: int = 5555, - backlog: int = 16, - ) -> None: - if host != "127.0.0.1": - raise ValueError("Python reference server only listens on 127.0.0.1") - if port < 0 or port > 65535: - raise ValueError("port must be in 0..65535") - self.application = ProtocolApplication(policy) - self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - self.socket.bind((host, port)) - self.socket.listen(backlog) - self.address = self.socket.getsockname() - self._closed = False - - def close(self) -> None: - if not self._closed: - self.socket.close() - self._closed = True - - def _send( - self, - client: socket.socket, - request: RequestHeader, - status: int, - payload: bytes, - ) -> None: - client.sendall( - _response_header(request, status=status, payload_len=len(payload)) - ) - if payload: - client.sendall(payload) - - def _handle_client(self, client: socket.socket) -> None: - while not self.application.shutdown_requested: - recv_started = time.perf_counter() - raw_header = _recv_exact( - client, - wire.HEADER_SIZE, - allow_initial_eof=True, - ) - if raw_header is None: - return - request = decode_request_header(raw_header) - if request.magic != wire.MAGIC: - return - if request.version != wire.VERSION: - self._send( - client, - request, - STATUS_BAD_VERSION, - b"bad protocol version", - ) - return - try: - validate_request_header(request) - except PayloadTooBig as exc: - self._send(client, request, STATUS_PAYLOAD_TOO_BIG, str(exc).encode("utf-8")) - return - except ProtocolError as exc: - self._send(client, request, STATUS_BAD_REQUEST, str(exc).encode("utf-8")) - return - payload = _recv_exact(client, request.payload_len) if request.payload_len else b"" - assert payload is not None - server_recv_ms = (time.perf_counter() - recv_started) * 1000.0 - try: - response, should_shutdown = self.application.dispatch( - request.op, - payload, - server_recv_ms=server_recv_ms, - ) - status = wire.STATUS_OK - except ProtocolError as exc: - response = str(exc).encode("utf-8") - should_shutdown = False - status = STATUS_BAD_REQUEST - except Exception as exc: - logging.exception("Python reference inference failed") - response = f"Python reference inference failed: {exc}".encode("utf-8") - should_shutdown = False - status = STATUS_INTERNAL_ERROR - self._send(client, request, status, response) - if should_shutdown: - return - - def serve_forever(self) -> None: - try: - while not self.application.shutdown_requested: - client, peer = self.socket.accept() - logging.debug("protocol connection from %s:%s", *peer) - with client: - try: - self._handle_client(client) - except (ConnectionError, OSError): - logging.debug("protocol peer disconnected", exc_info=True) - finally: - self.close() - - -def load_pinned_reference_policy( - *, - checkpoint_root: Path, - starvla_source: Path | None, - device: str, - noise_seed: int, - default_unnorm_key: str, - variant_name: str = "oft", -) -> PinnedOFTReferencePolicy: - source_dir = starvla_source or checkpoint_root / "source" / "starvla" - paths = validate_reference_inputs( - checkpoint_root=checkpoint_root, - source_dir=Path(source_dir), - variant_name=variant_name, - catalog_path=DEFAULT_CATALOG, - ) - try: - import torch - import transformers - except ImportError as exc: - raise StarVLAError(f"official StarVLA runtime dependency is missing: {exc}") from exc - validate_runtime_versions( - torch_version=torch.__version__, - torchvision_version=_distribution_version("torchvision"), - transformers_version=transformers.__version__, - numpy_version=np.__version__, - ) - if ( - variant_name == "qwen25_oft" - and _distribution_version("qwen-vl-utils") - != EXPECTED_QWEN_VL_UTILS_VERSION - ): - raise StarVLAError( - "official Qwen2.5 OFT reference requires qwen-vl-utils " - f"{EXPECTED_QWEN_VL_UTILS_VERSION}, got " - f"{_distribution_version('qwen-vl-utils')}" - ) - _configure_determinism(torch, seed=noise_seed, device=device) - framework_loader = ( - load_qwen25_official_framework - if variant_name == "qwen25_oft" - else load_qwen3_official_framework - ) - framework, _config = framework_loader(paths, device=device) - framework = apply_official_bf16(framework, torch) - - source_dir = Path(paths["source_dir"]) - processor_factory: Callable[..., Any] | None = None - action_unnormalizer: Callable[[np.ndarray, str], np.ndarray] | None = None - if variant_name == "qwen25_oft": - for profile in framework.norm_stats: - legacy_normalization_contract(framework.norm_stats, str(profile)) - - def qwen25_unnormalizer( - normalized: np.ndarray, profile: str - ) -> np.ndarray: - return unnormalize_legacy_actions( - normalized, - framework.norm_stats, - profile, - )[0] - - action_unnormalizer = qwen25_unnormalizer - else: - sys.path.insert(0, str(source_dir)) - try: - from deployment.model_server import policy_norm_processor - - _assert_module_origin(policy_norm_processor, source_dir) - processor_factory = policy_norm_processor.PolicyNormProcessor - finally: - if sys.path and sys.path[0] == str(source_dir): - del sys.path[0] - - metadata = build_server_metadata( - paths, - framework, - default_unnorm_key=default_unnorm_key, - source_tree_sha1=_git_tree_sha1(source_dir), - source_tracked_index_sha256=_git_tracked_index_sha256(source_dir), - runtime=build_runtime_metadata( - torch, - transformers, - device=device, - ), - ) - return PinnedOFTReferencePolicy( - framework=framework, - processor_factory=processor_factory, - checkpoint=Path(paths["checkpoint"]), - metadata=metadata, - action_unnormalizer=action_unnormalizer, - ) - - -def write_metadata(path: Path, metadata: Mapping[str, Any]) -> None: - path = path.resolve() - path.parent.mkdir(parents=True, exist_ok=True) - serialized = json.dumps(metadata, indent=2, sort_keys=True) + "\n" - with tempfile.NamedTemporaryFile( - mode="w", - encoding="utf-8", - dir=path.parent, - prefix=f".{path.name}.", - delete=False, - ) as handle: - temporary = Path(handle.name) - handle.write(serialized) - handle.flush() - os.fsync(handle.fileno()) - try: - os.replace(temporary, path) - finally: - temporary.unlink(missing_ok=True) - - -def build_parser() -> argparse.ArgumentParser: - parser = argparse.ArgumentParser( - description=( - "Serve a pinned official StarVLA Qwen-VL OFT checkpoint over " - "robot.cpp protocol v4." - ) - ) - parser.add_argument( - "--variant", - choices=SUPPORTED_VARIANTS, - default="oft", - help="Catalog variant; qwen25_oft uses the plain Qwen2.5-VL assets.", - ) - parser.add_argument("--checkpoint-root", type=Path, default=Path("ckpts/starvla")) - parser.add_argument("--starvla-source", type=Path) - parser.add_argument("--device", default="cuda:0") - parser.add_argument("--host", default="127.0.0.1") - parser.add_argument("--port", type=int, default=5555) - parser.add_argument( - "--unnorm-key", - help="Defaults to oxe_bridge for oft and bridge_dataset for qwen25_oft.", - ) - parser.add_argument( - "--preflight", - action="store_true", - help="Verify pinned local inputs and print provenance without loading the model.", - ) - parser.add_argument( - "--noise-seed", - type=int, - default=0, - help="Deterministic runtime seed; OFT inference itself has no sampled noise.", - ) - parser.add_argument( - "--metadata-output", - type=Path, - help="Optional atomic JSON record of all pinned identities and dtype contracts.", - ) - parser.add_argument("--verbosity", type=int, default=0) - return parser - - -def main(argv: Sequence[str] | None = None) -> int: - args = build_parser().parse_args(argv) - _require_isolated_python() - if args.host != "127.0.0.1": - raise StarVLAError("--host must be 127.0.0.1") - if args.port <= 0 or args.port > 65535: - raise StarVLAError("--port must be in 1..65535") - if args.verbosity < 0: - raise StarVLAError("--verbosity must be non-negative") - logging.basicConfig( - level=logging.DEBUG if args.verbosity else logging.INFO, - format="%(asctime)s %(levelname)s %(message)s", - force=True, - ) - default_unnorm_key = ( - args.unnorm_key or default_unnorm_key_for_variant(args.variant) - ) - - checkpoint_root = args.checkpoint_root.resolve() - source_dir = ( - args.starvla_source.resolve() - if args.starvla_source - else checkpoint_root / "source" / "starvla" - ) - if args.preflight: - paths = validate_reference_inputs( - checkpoint_root=checkpoint_root, - source_dir=source_dir, - variant_name=args.variant, - catalog_path=DEFAULT_CATALOG, - ) - record = build_preflight_record(paths) - serialized = json.dumps(record, indent=2, sort_keys=True) + "\n" - if args.metadata_output is not None: - write_metadata(args.metadata_output, record) - sys.stdout.write(serialized) - return 0 - - policy = load_pinned_reference_policy( - checkpoint_root=checkpoint_root, - starvla_source=source_dir, - device=args.device, - noise_seed=args.noise_seed, - default_unnorm_key=default_unnorm_key, - variant_name=args.variant, - ) - if args.metadata_output is not None: - write_metadata(args.metadata_output, policy.metadata) - logging.info( - "loaded pinned OFT Python reference metadata=%s", - _canonical_json_bytes(policy.metadata).decode("ascii"), - ) - - server = ReferenceProtocolServer(policy, host=args.host, port=args.port) - logging.info( - "Python reference server listening on %s:%d model=%s variant=%s", - server.address[0], - server.address[1], - MODEL_TYPE, - args.variant, - ) - try: - server.serve_forever() - except KeyboardInterrupt: - logging.info("Python reference server interrupted") - server.close() - return 0 - - -if __name__ == "__main__": - try: - raise SystemExit(main()) - except StarVLAError as exc: - raise SystemExit(f"error: {exc}") from exc diff --git a/tools/hf2gguf/starvla/starvla_checkpoint.py b/tools/hf2gguf/starvla/starvla_checkpoint.py index ef50d79..a9722ee 100755 --- a/tools/hf2gguf/starvla/starvla_checkpoint.py +++ b/tools/hf2gguf/starvla/starvla_checkpoint.py @@ -29,7 +29,7 @@ "qwen25_fast": "qwen25-fast", } -DEFAULT_QWEN_ASSET = "qwen3_vl_4b_instruct" +LEGACY_QWEN3_ASSET = "qwen3_vl_4b_instruct" SUPPORTED_BACKBONES = {"qwen3_vl", "qwen2_5_vl"} SUPPORTED_FRAMEWORKS = {"oft", "groot", "pi", "pi_v3", "fast"} GENERATED_QWEN_ASSET_PATHS = {"model.safetensors.index.json"} @@ -170,26 +170,29 @@ def load_catalog(path: Path | str = DEFAULT_CATALOG) -> dict[str, Any]: raise StarVLAError( f"catalog variant {name!r} has unsupported framework={framework!r}" ) - backbone = entry.get("backbone", "qwen3_vl") + if entry.get("model_type") != "starvla": + raise StarVLAError(f"catalog variant {name!r} must use model_type='starvla'") + backbone = entry.get("backbone") if backbone not in SUPPORTED_BACKBONES: raise StarVLAError( f"catalog variant {name!r} has unsupported backbone={backbone!r}" ) - qwen_asset = entry.get("qwen_asset", DEFAULT_QWEN_ASSET) + qwen_asset = entry.get("qwen_asset") if not isinstance(qwen_asset, str) or qwen_asset not in shared_assets: raise StarVLAError( f"catalog variant {name!r} references unknown qwen_asset={qwen_asset!r}" ) if not isinstance(entry.get("repo_id"), str) or not entry["repo_id"]: raise StarVLAError(f"catalog variant {name!r} is missing repo_id/revision") + if not isinstance(entry.get("default_unnorm_key"), str) or not entry["default_unnorm_key"]: + raise StarVLAError(f"catalog variant {name!r} has no default_unnorm_key") _validate_revision(entry.get("revision"), field=f"variant {name}.revision") checkpoint = entry.get("checkpoint") - if checkpoint is not None: - if not isinstance(checkpoint, dict): - raise StarVLAError(f"catalog variant {name!r} checkpoint must be an object or null") - _safe_relative_path(checkpoint.get("path"), field=f"variant {name}.checkpoint.path") - _validate_positive_size(checkpoint.get("size"), field=f"variant {name}.checkpoint.size") - _validate_sha256(checkpoint.get("sha256"), field=f"variant {name}.checkpoint.sha256") + if not isinstance(checkpoint, dict): + raise StarVLAError(f"catalog variant {name!r} checkpoint must be an object") + _safe_relative_path(checkpoint.get("path"), field=f"variant {name}.checkpoint.path") + _validate_positive_size(checkpoint.get("size"), field=f"variant {name}.checkpoint.size") + _validate_sha256(checkpoint.get("sha256"), field=f"variant {name}.checkpoint.sha256") entries = { **{f"shared asset {name}": entry for name, entry in shared_assets.items()}, **{f"variant {name}": entry for name, entry in variants.items()}, @@ -254,12 +257,80 @@ def get_variant(catalog: Mapping[str, Any], variant: str) -> dict[str, Any]: return entry +def local_checkpoint_catalog( + catalog: Mapping[str, Any], + variant_name: str, + checkpoint: Path, + source_dir: Path, + default_unnorm_key: str | None = None, +) -> dict[str, Any]: + """Bind a training checkpoint and its run metadata to a catalog variant.""" + if not checkpoint.is_file(): + raise StarVLAError(f"checkpoint does not exist: {checkpoint}") + required_assets = ("config.yaml", "dataset_statistics.json") + for name in required_assets: + if not (source_dir / name).is_file(): + raise StarVLAError(f"training run is missing {name}: {source_dir / name}") + try: + stats = json.loads((source_dir / "dataset_statistics.json").read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise StarVLAError(f"failed to load training dataset statistics: {exc}") from exc + if not isinstance(stats, dict) or not stats: + raise StarVLAError("training dataset_statistics.json must contain at least one profile") + + result = deepcopy(catalog) + variant = get_variant(result, variant_name) + selected_key = default_unnorm_key or str(variant["default_unnorm_key"]) + if selected_key not in stats: + if default_unnorm_key is None and len(stats) == 1: + selected_key = next(iter(stats)) + else: + raise StarVLAError( + f"normalization profile {selected_key!r} is not present; " + f"choose one of {sorted(stats)}" + ) + + local_entry = result["variants"][variant_name] + checkpoint_sha256 = sha256_file(checkpoint) + local_entry["repo_id"] = "local" + local_entry["revision"] = checkpoint_sha256[:40] + local_entry["checkpoint"] = { + "path": checkpoint.name, + "size": checkpoint.stat().st_size, + "sha256": checkpoint_sha256, + } + local_entry["files"] = list(required_assets) + local_entry["file_hashes"] = { + name: { + "size": (source_dir / name).stat().st_size, + "sha256": sha256_file(source_dir / name), + } + for name in required_assets + } + local_entry["default_unnorm_key"] = selected_key + return result + + +def portable_source_record( + source: Mapping[str, Any], variant_entry: Mapping[str, Any] +) -> dict[str, Any]: + """Replace the staging checkpoint path with its catalog-relative path.""" + checkpoint = variant_entry.get("checkpoint") + if not isinstance(checkpoint, Mapping): + raise StarVLAError("catalog variant has no policy checkpoint") + result = dict(source) + result["checkpoint"] = _safe_relative_path( + checkpoint.get("path"), field="variant checkpoint path" + ).as_posix() + return result + + def get_qwen_asset( catalog: Mapping[str, Any], variant_entry: Mapping[str, Any] ) -> tuple[str, dict[str, Any]]: - asset_name = str(variant_entry.get("qwen_asset", DEFAULT_QWEN_ASSET)) + asset_name = variant_entry.get("qwen_asset") shared_assets = catalog.get("shared_assets", {}) - if asset_name not in shared_assets: + if not isinstance(asset_name, str) or asset_name not in shared_assets: raise StarVLAError( f"variant {variant_entry.get('_catalog_key', variant_entry.get('framework'))!r} " f"references unknown Qwen asset {asset_name!r}" @@ -289,7 +360,7 @@ def sha256_file(path: Path, chunk_size: int = 8 * 1024 * 1024) -> str: def verify_checkpoint_file(path: Path, variant_entry: Mapping[str, Any]) -> None: checkpoint = variant_entry.get("checkpoint") if checkpoint is None: - raise StarVLAError(f"variant {variant_entry.get('framework')!r} has no official policy checkpoint") + raise StarVLAError(f"variant {variant_entry.get('framework')!r} has no policy checkpoint") if not path.is_file(): raise StarVLAError(f"checkpoint does not exist: {path}") expected_size = int(checkpoint["size"]) @@ -397,7 +468,7 @@ def validate_qwen_vlm_destination_names( ) -def official_bundle_uuid(variant_entry: Mapping[str, Any], catalog: Mapping[str, Any]) -> str: +def bundle_uuid(variant_entry: Mapping[str, Any], catalog: Mapping[str, Any]) -> str: """Derive the bundle identity from every source that can change runtime semantics.""" qwen_asset_name, qwen_entry = get_qwen_asset(catalog, variant_entry) qwen_hashes = staged_qwen_asset_hashes(qwen_entry) @@ -424,11 +495,11 @@ def official_bundle_uuid(variant_entry: Mapping[str, Any], catalog: Mapping[str, }, } catalog_variant = variant_entry.get("_catalog_key", variant_entry["framework"]) - backbone = variant_entry.get("backbone", "qwen3_vl") + backbone = variant_entry["backbone"] if ( catalog_variant != variant_entry["framework"] or backbone != "qwen3_vl" - or qwen_asset_name != DEFAULT_QWEN_ASSET + or qwen_asset_name != LEGACY_QWEN3_ASSET ): provenance["catalog_variant"] = catalog_variant provenance["backbone"] = backbone @@ -437,22 +508,10 @@ def official_bundle_uuid(variant_entry: Mapping[str, Any], catalog: Mapping[str, return str(uuid.uuid5(uuid.NAMESPACE_URL, f"robotcpp:starvla-bundle:{canonical}")) -def _flatten_config(value: Any, prefix: str = "") -> dict[str, Any]: - if not isinstance(value, Mapping): - return {prefix: value} - flattened: dict[str, Any] = {} - for key, child in value.items(): - path = f"{prefix}.{key}" if prefix else str(key) - flattened.update(_flatten_config(child, path)) - return flattened - - def _set_effective_value( effective: dict[str, Any], path: str, value: Any, - authority: str, - overrides: dict[str, dict[str, Any]], ) -> None: owner: dict[str, Any] = effective parts = path.split(".") @@ -464,10 +523,7 @@ def _set_effective_value( if not isinstance(child, dict): raise StarVLAError(f"effective config path is not an object: {path}") owner = child - previous = owner.get(parts[-1]) owner[parts[-1]] = value - if previous != value: - overrides[path] = {"source": previous, "effective": value, "authority": authority} def resolve_effective_config( @@ -480,7 +536,7 @@ def resolve_effective_config( str(variant_entry["framework"]) if variant_entry is not None else variant_name ) backbone = ( - str(variant_entry.get("backbone", "qwen3_vl")) + str(variant_entry["backbone"]) if variant_entry is not None else "qwen3_vl" ) @@ -505,109 +561,24 @@ def resolve_effective_config( if not isinstance(canonical, dict): raise StarVLAError(f"expected an object in canonical StarVLA config {yaml_path}") - candidate_conflicts: dict[str, dict[str, Any]] = {} - json_path = source_dir / "config.json" - if json_path.is_file(): - try: - json_config = json.loads(json_path.read_text(encoding="utf-8")) - except (OSError, json.JSONDecodeError) as exc: - raise StarVLAError(f"failed to load StarVLA config mirror {json_path}: {exc}") from exc - canonical_flat = _flatten_config(canonical) - json_flat = _flatten_config(json_config) - if set(json_flat) != set(canonical_flat): - raise StarVLAError(f"config.json and canonical config.yaml have different keys in {source_dir}") - conflicts = { - path: {"config.yaml": canonical_flat[path], "config.json": json_flat[path]} - for path in canonical_flat - if canonical_flat[path] != json_flat[path] - } - allowed = {"framework.qwenvl.base_vlm"} if framework_name == "groot" else set() - unexpected = set(conflicts) - allowed - if unexpected: - raise StarVLAError( - f"config.json and canonical config.yaml disagree at unsupported paths in {source_dir}: " - f"{sorted(unexpected)}" - ) - candidate_conflicts.update(conflicts) - elif backbone == "qwen3_vl" and framework_name in {"oft", "groot"}: - raise StarVLAError(f"missing StarVLA config mirror: {json_path}") - - if framework_name == "pi_v3": - full_path = source_dir / "config.full.yaml" - if not full_path.is_file(): - raise StarVLAError(f"missing PI_v3 full config candidate: {full_path}") - try: - full_config = yaml.safe_load(full_path.read_text(encoding="utf-8")) - except (OSError, yaml.YAMLError) as exc: - raise StarVLAError(f"failed to load PI_v3 full config {full_path}: {exc}") from exc - if not isinstance(full_config, dict): - raise StarVLAError(f"expected an object in PI_v3 full config {full_path}") - canonical_flat = _flatten_config(canonical) - full_flat = _flatten_config(full_config) - conflicts = { - path: {"config.yaml": canonical_flat[path], "config.full.yaml": full_flat[path]} - for path in set(canonical_flat) & set(full_flat) - if canonical_flat[path] != full_flat[path] - } - allowed = {"framework.action_model.diffusion_model_cfg.interleave_self_attention"} - unexpected = set(conflicts) - allowed - if unexpected: - raise StarVLAError( - f"PI_v3 config.full.yaml disagrees with canonical config.yaml at unsupported paths: " - f"{sorted(unexpected)}" - ) - candidate_conflicts.update(conflicts) - effective = deepcopy(canonical) - overrides: dict[str, dict[str, Any]] = {} qwen_hidden_dim = 2048 if backbone == "qwen2_5_vl" else 2560 if framework_name == "oft": - _set_effective_value( - effective, - "framework.qwenvl.vl_hidden_dim", - qwen_hidden_dim, - "checkpoint_tensor_shape", - overrides, - ) - _set_effective_value( - effective, - "framework.action_model.action_hidden_dim", - qwen_hidden_dim, - "checkpoint_tensor_shape", - overrides, - ) - _set_effective_value( - effective, - "framework.action_model.action_model_type", - "MLP", - "pinned_starvla_qwenoft_factory_and_checkpoint_topology", - overrides, + values = ( + ("framework.qwenvl.vl_hidden_dim", qwen_hidden_dim), + ("framework.action_model.action_hidden_dim", qwen_hidden_dim), + ("framework.action_model.action_model_type", "MLP"), ) elif framework_name == "groot": - _set_effective_value( - effective, - "framework.qwenvl.vl_hidden_dim", - qwen_hidden_dim, - "checkpoint_tensor_shape", - overrides, - ) - _set_effective_value( - effective, - "framework.action_model.diffusion_model_cfg.cross_attention_dim", - qwen_hidden_dim, - "pinned_starvla_qwengroot_runtime_and_checkpoint_tensor_shape", - overrides, - ) - for path, value in ( + values = ( + ("framework.qwenvl.vl_hidden_dim", qwen_hidden_dim), + ("framework.action_model.diffusion_model_cfg.cross_attention_dim", qwen_hidden_dim), ("framework.action_model.diffusion_model_cfg.input_embedding_dim", 768), ("framework.action_model.diffusion_model_cfg.attention_head_dim", 64), ("framework.action_model.diffusion_model_cfg.num_attention_heads", 12), - ): - _set_effective_value( - effective, path, value, "pinned_starvla_dit_b_definition", overrides - ) + ) elif framework_name == "pi_v3": - for path, value in ( + values = ( ("framework.qwenvl.vl_hidden_dim", 2560), ("framework.qwenvl.num_vl_layers", 36), ("framework.action_model.action_model_type", "LayerwiseFM"), @@ -619,16 +590,9 @@ def resolve_effective_config( ("framework.action_model.diffusion_model_cfg.num_layers", 36), ("framework.action_model.diffusion_model_cfg.interleave_self_attention", False), ("framework.action_model.diffusion_model_cfg.use_canonical_forward", True), - ): - _set_effective_value( - effective, - path, - value, - "pinned_starvla_qwenpi_v3_runtime_and_released_checkpoint_config", - overrides, - ) + ) elif framework_name == "pi": - for path, value in ( + values = ( ("framework.qwenvl.vl_hidden_dim", qwen_hidden_dim), ("framework.action_model.hidden_size", qwen_hidden_dim), ( @@ -645,38 +609,18 @@ def resolve_effective_config( qwen_hidden_dim // 64, ), ("framework.action_model.diffusion_model_cfg.use_canonical_forward", False), - ): - _set_effective_value( - effective, - path, - value, - "pinned_starvla_qwenpi_runtime_and_checkpoint_tensor_shape", - overrides, - ) + ) else: raise AssertionError(f"unhandled effective-config framework: {framework_name}") - _set_effective_value(effective, "framework.action_model.action_horizon", 16, "released_checkpoint_contract", overrides) - _set_effective_value(effective, "version_id", "0.21", "pinned_starvla_config_compat", overrides) - - inactive_fields = [] - if framework_name == "oft": - inactive_fields = [ - "framework.action_model.diffusion_model_cfg", - "framework.action_model.hidden_size", - "framework.action_model.state_dim", - ] - elif framework_name == "groot": - inactive_fields = ["framework.action_model.action_hidden_dim"] + for path, value in values: + _set_effective_value(effective, path, value) + _set_effective_value(effective, "framework.action_model.action_horizon", 16) effective["_robotcpp_effective_config"] = { "schema_version": 1, "variant": variant_name, "framework": framework_name, "backbone": backbone, - "canonical_source": "config.yaml", - "candidate_conflicts": candidate_conflicts, - "overrides": overrides, - "inactive_fields": inactive_fields, } return effective @@ -688,9 +632,14 @@ def load_checkpoint_state(path: Path) -> dict[str, Any]: raise StarVLAError("PyTorch is required to inspect a StarVLA checkpoint") from exc try: - raw = torch.load(path, map_location="cpu", mmap=True, weights_only=True) + if path.suffix == ".safetensors": + from safetensors.torch import load_file + + raw = load_file(path, device="cpu") + else: + raw = torch.load(path, map_location="cpu", mmap=True, weights_only=True) except Exception as exc: - raise StarVLAError(f"failed to load checkpoint {path} with weights_only=True: {exc}") from exc + raise StarVLAError(f"failed to load checkpoint {path}: {exc}") from exc if isinstance(raw, Mapping) and raw and all(isinstance(key, str) and torch.is_tensor(value) for key, value in raw.items()): return dict(raw) @@ -712,7 +661,7 @@ def load_checkpoint_state(path: Path) -> dict[str, Any]: def classify_tensor(name: str, variant_entry: Mapping[str, Any]) -> tuple[str, str, str]: - backbone = str(variant_entry.get("backbone", "qwen3_vl")) + backbone = str(variant_entry["backbone"]) destination_prefixes = VLM_DESTINATION_PREFIXES.get(backbone) if destination_prefixes is None: raise StarVLAError(f"unsupported StarVLA Qwen backbone: {backbone!r}") @@ -848,15 +797,15 @@ def validate_expected_inventory(records: list[TensorRecord], variant_entry: Mapp raise StarVLAError("checkpoint required-shape mismatch: " + "; ".join(shape_mismatches)) -def validate_official_surgery_manifest( +def validate_surgery_manifest( manifest: Mapping[str, Any], variant_entry: Mapping[str, Any], catalog: Mapping[str, Any], ) -> None: - """Require a surgery manifest to describe the pinned official checkpoint exactly.""" + """Validate a surgery manifest against its catalog entry.""" checkpoint = variant_entry.get("checkpoint") if checkpoint is None: - raise StarVLAError(f"variant {variant_entry.get('framework')!r} has no official checkpoint") + raise StarVLAError(f"variant {variant_entry.get('framework')!r} has no checkpoint") expected_top_level = { "schema_version": 1, @@ -893,7 +842,7 @@ def validate_official_surgery_manifest( if inventory.get(key) != expected: mismatches.append(f"inventory.{key}: expected {expected!r}, got {inventory.get(key)!r}") - expected_uuid = official_bundle_uuid(variant_entry, catalog) + expected_uuid = bundle_uuid(variant_entry, catalog) if manifest.get("bundle_uuid") != expected_uuid: mismatches.append(f"bundle_uuid: expected {expected_uuid!r}, got {manifest.get('bundle_uuid')!r}") @@ -946,7 +895,7 @@ def validate_official_surgery_manifest( ): mismatches.append("effective_config: invalid path/size/SHA256 record") if mismatches: - raise StarVLAError("non-official or inconsistent surgery manifest: " + "; ".join(mismatches)) + raise StarVLAError("surgery manifest does not match the catalog: " + "; ".join(mismatches)) def verify_staged_assets(root: Path, assets: Mapping[str, Any], *, component: str) -> None: @@ -1032,7 +981,7 @@ def _verify_staged_component( expected_names = {record.destination_name for record in records} if set(weight_map) != expected_names: raise StarVLAError( - f"staged {component} tensor set does not match the official checkpoint: " + f"staged {component} tensor set does not match the checkpoint: " f"expected {len(expected_names)}, got {len(weight_map)}" ) @@ -1062,7 +1011,7 @@ def _verify_staged_component( ) if not torch.equal(staged, original): raise StarVLAError( - f"staged tensor content does not match the official checkpoint: {record.destination_name}" + f"staged tensor content does not match the checkpoint: {record.destination_name}" ) del staged @@ -1086,7 +1035,7 @@ def verify_staged_components_against_checkpoint( manifest_records = manifest.get("tensors") expected_manifest = [record.to_json() for record in source_records] if manifest_records != expected_manifest: - raise StarVLAError("surgery tensor inventory does not match the verified official checkpoint") + raise StarVLAError("surgery tensor inventory does not match the verified checkpoint") for component, (root, output) in components.items(): _verify_staged_component( root, diff --git a/tools/hf2gguf/starvla/starvla_surgery.py b/tools/hf2gguf/starvla/starvla_surgery.py index 54831ed..d5d23b3 100755 --- a/tools/hf2gguf/starvla/starvla_surgery.py +++ b/tools/hf2gguf/starvla/starvla_surgery.py @@ -23,7 +23,7 @@ inventory_summary, load_catalog, load_checkpoint_state, - official_bundle_uuid, + bundle_uuid, resolve_effective_config, sha256_file, staged_qwen_asset_hashes, @@ -194,21 +194,17 @@ def _run_surgery_in_owned_directory( variant_name: str, catalog_path: Path, max_shard_size: int, - *, - verify_hash: bool, - enforce_expected: bool, ) -> dict[str, Any]: catalog = load_catalog(catalog_path) variant = get_variant(catalog, variant_name) if variant.get("checkpoint") is None: - raise StarVLAError(f"variant {variant_name!r} has no official policy checkpoint to split") - if verify_hash: - verify_checkpoint_file(checkpoint, variant) + raise StarVLAError(f"variant {variant_name!r} has no policy checkpoint to split") + verify_checkpoint_file(checkpoint, variant) if output_dir.exists() and any(output_dir.iterdir()): raise StarVLAError(f"output directory is not empty: {output_dir}") state_dict = load_checkpoint_state(checkpoint) - records = build_inventory(state_dict, variant, enforce_expected=enforce_expected) + records = build_inventory(state_dict, variant, enforce_expected=True) vlm_records = [record for record in records if record.component == "vlm"] policy_records = [record for record in records if record.component == "policy"] if len(vlm_records) + len(policy_records) != len(records): @@ -224,7 +220,7 @@ def _run_surgery_in_owned_directory( base_assets, qwen_asset_entry, vlm_records, - backbone=str(variant.get("backbone", "qwen3_vl")), + backbone=str(variant["backbone"]), ) qwen_assets = copy_qwen_assets(base_assets, hf_dir, qwen_asset_entry) policy_assets = copy_policy_assets(source_dir, policy_dir, variant) @@ -255,21 +251,20 @@ def _run_surgery_in_owned_directory( max_shard_size, ) - checkpoint_sha256 = sha256_file(checkpoint) if not verify_hash else str(variant["checkpoint"]["sha256"]) - bundle_uuid = official_bundle_uuid(variant, catalog) + source_uuid = bundle_uuid(variant, catalog) manifest = { "schema_version": 1, "variant": variant_name, "framework": variant["framework"], - "backbone": variant.get("backbone", "qwen3_vl"), + "backbone": variant["backbone"], "model_type": variant["model_type"], - "bundle_uuid": bundle_uuid, + "bundle_uuid": source_uuid, "source": { "repo_id": variant["repo_id"], "revision": variant["revision"], "checkpoint": str(checkpoint), "checkpoint_size": checkpoint.stat().st_size, - "checkpoint_sha256": checkpoint_sha256, + "checkpoint_sha256": variant["checkpoint"]["sha256"], "starvla_revision": catalog["source_revisions"]["starvla"], "llama_cpp_revision": catalog["source_revisions"]["llama_cpp"], "qwen_repo_id": qwen_asset_entry["repo_id"], @@ -296,9 +291,6 @@ def run_surgery( variant_name: str, catalog_path: Path, max_shard_size: int, - *, - verify_hash: bool, - enforce_expected: bool, ) -> dict[str, Any]: """Own the staging directory so a failed split cannot poison a retry.""" output_dir.parent.mkdir(parents=True, exist_ok=True) @@ -316,8 +308,6 @@ def run_surgery( variant_name=variant_name, catalog_path=catalog_path, max_shard_size=max_shard_size, - verify_hash=verify_hash, - enforce_expected=enforce_expected, ) except BaseException: shutil.rmtree(output_dir, ignore_errors=True) @@ -337,8 +327,6 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--output-dir", type=Path, required=True) parser.add_argument("--catalog", type=Path, default=DEFAULT_CATALOG) parser.add_argument("--max-shard-size", type=parse_size, default=parse_size("2G")) - parser.add_argument("--skip-hash-check", action="store_true") - parser.add_argument("--allow-nonofficial-inventory", action="store_true") return parser.parse_args() @@ -353,8 +341,6 @@ def main() -> int: variant_name=args.variant, catalog_path=args.catalog, max_shard_size=args.max_shard_size, - verify_hash=not args.skip_hash_check, - enforce_expected=not args.allow_nonofficial_inventory, ) print(json.dumps(manifest["inventory"], indent=2, sort_keys=True)) print(f"surgery manifest: {args.output_dir / 'surgery_manifest.json'}") diff --git a/tools/hf2gguf/starvla/starvla_variant_config.sh b/tools/hf2gguf/starvla/starvla_variant_config.sh new file mode 100644 index 0000000..a4b8076 --- /dev/null +++ b/tools/hf2gguf/starvla/starvla_variant_config.sh @@ -0,0 +1,55 @@ +#!/usr/bin/env bash + +_STARVLA_CONFIG_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" + +load_starvla_variant() { + if [[ $# -ne 1 ]]; then + echo "usage: load_starvla_variant VARIANT" >&2 + return 2 + fi + + local config_python="${STARVLA_CONFIG_PYTHON:-python3}" + local assignments + assignments="$("${config_python}" - "${_STARVLA_CONFIG_DIR}" "$1" <<'PY' +import os +import shlex +import sys +from pathlib import Path + +sys.path.insert(0, sys.argv[1]) +from starvla_checkpoint import ( # noqa: E402 + StarVLAError, + artifact_stem, + get_qwen_asset, + get_variant, + load_catalog, +) + +try: + catalog_path = Path(os.environ["STARVLA_CATALOG"]) if "STARVLA_CATALOG" in os.environ else None + catalog = load_catalog(catalog_path) if catalog_path is not None else load_catalog() + variant = get_variant(catalog, sys.argv[2]) + _, qwen = get_qwen_asset(catalog, variant) +except StarVLAError as exc: + print(f"error: {exc}", file=sys.stderr) + raise SystemExit(2) from exc + +checkpoint = variant["checkpoint"] +values = { + "STARVLA_REVISION": catalog["source_revisions"]["starvla"], + "MODEL_TYPE": variant["model_type"], + "FRAMEWORK": variant["framework"], + "CHECKPOINT_REVISION": variant["revision"], + "CHECKPOINT_SHA256": checkpoint["sha256"], + "CHECKPOINT_DIRECTORY": variant["directory"], + "CHECKPOINT_RELATIVE_PATH": checkpoint["path"], + "QWEN_REVISION": qwen["revision"], + "QWEN_DIRECTORY": qwen["directory"], + "ARTIFACT_STEM": artifact_stem(sys.argv[2]), +} +for name, value in values.items(): + print(f"{name}={shlex.quote(str(value))}") +PY + )" || return 2 + eval "${assignments}" +} diff --git a/tools/hf2gguf/starvla/validate_starvla_bundle.py b/tools/hf2gguf/starvla/validate_starvla_bundle.py index 560b27a..6cf0a80 100755 --- a/tools/hf2gguf/starvla/validate_starvla_bundle.py +++ b/tools/hf2gguf/starvla/validate_starvla_bundle.py @@ -4,7 +4,6 @@ from __future__ import annotations import argparse -import hashlib import json import math import sys @@ -15,23 +14,14 @@ import numpy as np from convert_starvla_policy_to_gguf import ( - GROOT_BLOCK_COUNT, - GROOT_OFFICIAL_DIMENSIONS, - GROOT_OFFICIAL_DIMENSIONS_BY_BACKBONE, - GROOT_POLICY_TENSOR_COUNT, + GROOT_SUPPORTED_DIMENSIONS_BY_BACKBONE, GROOT_TENSOR_MAP, OFT_ACTION_TOKEN_ID, OFT_TENSOR_MAP, - PI_BLOCK_COUNT, - PI_OFFICIAL_DIMENSIONS, - PI_POLICY_TENSOR_COUNT, + PI_SUPPORTED_DIMENSIONS, PI_TENSOR_MAP, - PI_V3_BLOCK_COUNT, - PI_V3_OFFICIAL_DIMENSIONS, - PI_V3_POLICY_TENSOR_COUNT, - PI_V3_PROJECTOR_COUNT, + PI_V3_SUPPORTED_DIMENSIONS, PI_V3_TENSOR_MAP, - QWEN3VL_DYNAMIC_IMAGE_METADATA, build_groot_metadata, build_oft_metadata, build_pi_metadata, @@ -48,8 +38,9 @@ atomic_write_json, get_variant, load_catalog, + portable_source_record, sha256_file, - validate_official_surgery_manifest, + validate_surgery_manifest, verify_staged_assets, verify_staged_components_against_checkpoint, ) @@ -266,174 +257,6 @@ def expected_mmproj_tensor_map( return expected -def expected_groot_policy_tensor_map( - qwen_hidden_dim: int = 2560, -) -> dict[str, list[int]]: - """Return all GR00T GGUF shapes in ggml `ne[]` dimension order.""" - expected = { - "starvla.policy.groot.timestep.input.weight": [256, 768], - "starvla.policy.groot.timestep.input.bias": [768], - "starvla.policy.groot.timestep.output.weight": [768, 768], - "starvla.policy.groot.timestep.output.bias": [768], - } - for block in range(GROOT_BLOCK_COUNT): - attention_input_dim = qwen_hidden_dim if block % 2 == 0 else 768 - prefix = f"starvla.policy.groot.block.{block}" - expected.update( - { - f"{prefix}.ada_norm.weight": [768, 1536], - f"{prefix}.ada_norm.bias": [1536], - f"{prefix}.attention.query.weight": [768, 768], - f"{prefix}.attention.query.bias": [768], - f"{prefix}.attention.key.weight": [attention_input_dim, 768], - f"{prefix}.attention.key.bias": [768], - f"{prefix}.attention.value.weight": [attention_input_dim, 768], - f"{prefix}.attention.value.bias": [768], - f"{prefix}.attention.output.weight": [768, 768], - f"{prefix}.attention.output.bias": [768], - f"{prefix}.feed_forward.input.weight": [768, 3072], - f"{prefix}.feed_forward.input.bias": [3072], - f"{prefix}.feed_forward.output.weight": [3072, 768], - f"{prefix}.feed_forward.output.bias": [768], - } - ) - expected.update( - { - "starvla.policy.groot.output.modulation.weight": [768, 1536], - "starvla.policy.groot.output.modulation.bias": [1536], - "starvla.policy.groot.output.projection.weight": [768, 1024], - "starvla.policy.groot.output.projection.bias": [1024], - "starvla.policy.groot.action.input.weight": [7, 768], - "starvla.policy.groot.action.input.bias": [768], - "starvla.policy.groot.action.time_mix.weight": [1536, 768], - "starvla.policy.groot.action.time_mix.bias": [768], - "starvla.policy.groot.action.output.weight": [768, 768], - "starvla.policy.groot.action.output.bias": [768], - "starvla.policy.groot.velocity.input.weight": [1024, 1024], - "starvla.policy.groot.velocity.input.bias": [1024], - "starvla.policy.groot.velocity.output.weight": [1024, 7], - "starvla.policy.groot.velocity.output.bias": [7], - "starvla.policy.groot.future_tokens.weight": [768, 32], - "starvla.policy.groot.action_position.weight": [768, 1024], - } - ) - if len(expected) != GROOT_POLICY_TENSOR_COUNT: - raise AssertionError(f"internal GR00T tensor contract has {len(expected)} tensors") - return expected - - -def expected_pi_policy_tensor_map() -> dict[str, list[int]]: - """Return all legacy PI GGUF shapes in ggml `ne[]` dimension order.""" - expected = { - "starvla.policy.pi.timestep.input.weight": [256, 2048], - "starvla.policy.pi.timestep.input.bias": [2048], - "starvla.policy.pi.timestep.output.weight": [2048, 2048], - "starvla.policy.pi.timestep.output.bias": [2048], - } - for block in range(PI_BLOCK_COUNT): - prefix = f"starvla.policy.pi.block.{block}" - expected.update( - { - f"{prefix}.ada_norm.weight": [2048, 4096], - f"{prefix}.ada_norm.bias": [4096], - f"{prefix}.attention.query.weight": [2048, 2048], - f"{prefix}.attention.query.bias": [2048], - f"{prefix}.attention.key.weight": [2048, 2048], - f"{prefix}.attention.key.bias": [2048], - f"{prefix}.attention.value.weight": [2048, 2048], - f"{prefix}.attention.value.bias": [2048], - f"{prefix}.attention.output.weight": [2048, 2048], - f"{prefix}.attention.output.bias": [2048], - f"{prefix}.feed_forward.input.weight": [2048, 8192], - f"{prefix}.feed_forward.input.bias": [8192], - f"{prefix}.feed_forward.output.weight": [8192, 2048], - f"{prefix}.feed_forward.output.bias": [2048], - } - ) - expected.update( - { - "starvla.policy.pi.state.input.weight": [7, 2048], - "starvla.policy.pi.state.input.bias": [2048], - "starvla.policy.pi.state.output.weight": [2048, 2048], - "starvla.policy.pi.state.output.bias": [2048], - "starvla.policy.pi.action.input.weight": [7, 2048], - "starvla.policy.pi.action.input.bias": [2048], - "starvla.policy.pi.action.time_mix.weight": [4096, 2048], - "starvla.policy.pi.action.time_mix.bias": [2048], - "starvla.policy.pi.action.output.weight": [2048, 2048], - "starvla.policy.pi.action.output.bias": [2048], - "starvla.policy.pi.velocity.input.weight": [2048, 2048], - "starvla.policy.pi.velocity.input.bias": [2048], - "starvla.policy.pi.velocity.output.weight": [2048, 7], - "starvla.policy.pi.velocity.output.bias": [7], - "starvla.policy.pi.future_tokens.weight": [2048, 32], - "starvla.policy.pi.action_position.weight": [2048, 1024], - } - ) - if len(expected) != PI_POLICY_TENSOR_COUNT: - raise AssertionError(f"internal legacy PI tensor contract has {len(expected)} tensors") - return expected - - -def expected_pi_v3_policy_tensor_map() -> dict[str, list[int]]: - """Return all PI_v3 GGUF shapes in ggml `ne[]` dimension order.""" - expected = { - "starvla.policy.pi_v3.timestep.input.weight": [256, 1024], - "starvla.policy.pi_v3.timestep.input.bias": [1024], - "starvla.policy.pi_v3.timestep.output.weight": [1024, 1024], - "starvla.policy.pi_v3.timestep.output.bias": [1024], - } - for block in range(PI_V3_BLOCK_COUNT): - prefix = f"starvla.policy.pi_v3.block.{block}" - expected.update( - { - f"{prefix}.ada_norm.weight": [1024, 2048], - f"{prefix}.ada_norm.bias": [2048], - f"{prefix}.attention.query.weight": [1024, 1024], - f"{prefix}.attention.query.bias": [1024], - f"{prefix}.attention.key.weight": [1024, 1024], - f"{prefix}.attention.key.bias": [1024], - f"{prefix}.attention.value.weight": [1024, 1024], - f"{prefix}.attention.value.bias": [1024], - f"{prefix}.attention.output.weight": [1024, 1024], - f"{prefix}.attention.output.bias": [1024], - f"{prefix}.feed_forward.input.weight": [1024, 4096], - f"{prefix}.feed_forward.input.bias": [4096], - f"{prefix}.feed_forward.output.weight": [4096, 1024], - f"{prefix}.feed_forward.output.bias": [1024], - } - ) - expected.update( - { - "starvla.policy.pi_v3.action.input.weight": [7, 1024], - "starvla.policy.pi_v3.action.input.bias": [1024], - "starvla.policy.pi_v3.action.time_mix.weight": [2048, 1024], - "starvla.policy.pi_v3.action.time_mix.bias": [1024], - "starvla.policy.pi_v3.action.output.weight": [1024, 1024], - "starvla.policy.pi_v3.action.output.bias": [1024], - "starvla.policy.pi_v3.velocity.input.weight": [1024, 1024], - "starvla.policy.pi_v3.velocity.input.bias": [1024], - "starvla.policy.pi_v3.velocity.output.weight": [1024, 7], - "starvla.policy.pi_v3.velocity.output.bias": [7], - "starvla.policy.pi_v3.future_tokens.weight": [1024, 32], - "starvla.policy.pi_v3.action_position.weight": [1024, 1024], - } - ) - for projector in range(PI_V3_PROJECTOR_COUNT): - prefix = f"starvla.policy.pi_v3.projector.{projector}" - expected.update( - { - f"{prefix}.norm.weight": [2560], - f"{prefix}.norm.bias": [2560], - f"{prefix}.projection.weight": [2560, 1024], - f"{prefix}.projection.bias": [1024], - } - ) - if len(expected) != PI_V3_POLICY_TENSOR_COUNT: - raise AssertionError(f"internal PI_v3 tensor contract has {len(expected)} tensors") - return expected - - def metadata_matches(actual: Any, expected: Any) -> bool: if isinstance(expected, float): return isinstance(actual, (int, float)) and math.isclose(actual, expected, rel_tol=1e-6, abs_tol=1e-6) @@ -635,20 +458,6 @@ def validate_policy_metadata(reader: Any, expected: dict[str, Any]) -> None: expect_field(reader, "general.name", expected["general.name"]) -def validate_qwen3vl_image_metadata(reader: Any) -> None: - """Reject fixed-size or incomplete substitutes for the dynamic image contract.""" - expected_keys = set(QWEN3VL_DYNAMIC_IMAGE_METADATA) - actual_keys = {key for key in reader.fields if key.startswith("starvla.image.")} - if actual_keys != expected_keys: - raise StarVLAError( - "Qwen3-VL dynamic image metadata set mismatch; " - f"missing={sorted(expected_keys - actual_keys)}, " - f"unexpected={sorted(actual_keys - expected_keys)}" - ) - for key, expected in sorted(QWEN3VL_DYNAMIC_IMAGE_METADATA.items()): - expect_metadata_field(reader, key, expected) - - def validate_qwen_vl_image_metadata( reader: Any, expected_metadata: dict[str, Any], backbone: str ) -> None: @@ -738,8 +547,17 @@ def validate_policy_tensor_bytes( f"{len(tensor_name_map)}-tensor map" ) for source_name, destination_name in tensor_name_map.items(): - expected = _convert_policy_tensor_data(source_tensors[source_name], dtype) - actual = np.asarray(tensors[destination_name].data) + source = source_tensors[source_name] + expected_shape = list(reversed(source.shape)) + tensor = tensors[destination_name] + actual_shape = [int(dimension) for dimension in tensor.shape] + if actual_shape != expected_shape: + raise StarVLAError( + f"policy GGUF tensor shape mismatch for {destination_name}: " + f"expected {expected_shape}, got {actual_shape}" + ) + expected = _convert_policy_tensor_data(source, dtype) + actual = np.asarray(tensor.data) if actual.nbytes != expected.nbytes: raise StarVLAError( f"policy GGUF tensor byte size mismatch for {destination_name}: " @@ -978,40 +796,12 @@ def validate_policy( ) validate_policy_metadata(reader, expected_metadata) tensors = tensor_map(reader) - if framework == "oft": - tensor_name_map = OFT_TENSOR_MAP - input_dim = 2048 if backbone == "qwen2_5_vl" else 2560 - hidden_dim = 4096 if backbone == "qwen2_5_vl" else 5120 - expected_shapes = { - "starvla.policy.oft.input_norm.weight": [input_dim], - "starvla.policy.oft.input_norm.bias": [input_dim], - "starvla.policy.oft.input_proj.weight": [input_dim, hidden_dim], - "starvla.policy.oft.input_proj.bias": [hidden_dim], - "starvla.policy.oft.output_norm.weight": [hidden_dim], - "starvla.policy.oft.output_norm.bias": [hidden_dim], - "starvla.policy.oft.output_proj.weight": [hidden_dim, 7], - "starvla.policy.oft.output_proj.bias": [7], - } - for block in (0, 1): - expected_shapes[f"starvla.policy.oft.block.{block}.norm.weight"] = [hidden_dim] - expected_shapes[f"starvla.policy.oft.block.{block}.norm.bias"] = [hidden_dim] - expected_shapes[f"starvla.policy.oft.block.{block}.linear.weight"] = [ - hidden_dim, - hidden_dim, - ] - expected_shapes[f"starvla.policy.oft.block.{block}.linear.bias"] = [hidden_dim] - elif framework == "groot": - tensor_name_map = GROOT_TENSOR_MAP - expected_shapes = expected_groot_policy_tensor_map( - 2048 if backbone == "qwen2_5_vl" else 2560 - ) - elif framework == "pi": - tensor_name_map = PI_TENSOR_MAP - expected_shapes = expected_pi_policy_tensor_map() - else: - tensor_name_map = PI_V3_TENSOR_MAP - expected_shapes = expected_pi_v3_policy_tensor_map() - expect_complete_tensor_map(tensors, expected_shapes, f"{framework.upper()} policy") + tensor_name_map = { + "oft": OFT_TENSOR_MAP, + "groot": GROOT_TENSOR_MAP, + "pi": PI_TENSOR_MAP, + "pi_v3": PI_V3_TENSOR_MAP, + }[framework] dtype_counts = validate_dtype_set(reader, dtype, component="policy", exact=True) validate_policy_tensor_bytes( tensors, @@ -1088,8 +878,8 @@ def main() -> int: catalog = load_catalog(args.catalog) variant = get_variant(catalog, args.variant) framework = str(variant["framework"]) - backbone = str(variant.get("backbone", "qwen3_vl")) - validate_official_surgery_manifest(surgery_manifest, variant, catalog) + backbone = str(variant["backbone"]) + validate_surgery_manifest(surgery_manifest, variant, catalog) verify_staged_assets(args.hf_dir, surgery_manifest.get("qwen_assets", {}), component="Qwen") verify_staged_assets(args.policy_dir, surgery_manifest.get("policy_assets", {}), component="policy") verify_staged_components_against_checkpoint( @@ -1119,7 +909,7 @@ def main() -> int: ) elif framework == "groot": groot_dimensions = dict( - GROOT_OFFICIAL_DIMENSIONS_BY_BACKBONE[backbone] + GROOT_SUPPORTED_DIMENSIONS_BY_BACKBONE[backbone] ) expected_policy_metadata = build_groot_metadata( args.policy_dir, @@ -1136,7 +926,7 @@ def main() -> int: args.hf_dir, variant, surgery_manifest, - dict(PI_OFFICIAL_DIMENSIONS), + dict(PI_SUPPORTED_DIMENSIONS), args.text.name, args.mmproj.name, ) @@ -1146,7 +936,7 @@ def main() -> int: args.hf_dir, variant, surgery_manifest, - dict(PI_V3_OFFICIAL_DIMENSIONS), + dict(PI_V3_SUPPORTED_DIMENSIONS), args.text.name, args.mmproj.name, ) @@ -1199,13 +989,8 @@ def main() -> int: "variant": args.variant, "model_type": variant["model_type"], "bundle_uuid": bundle_uuid, - "source": surgery_manifest["source"], + "source": portable_source_record(surgery_manifest["source"], variant), "source_tensor_roles": dict(sorted(role_counts.items())), - "surgery_manifest": { - "filename": args.surgery_manifest.name, - "size": args.surgery_manifest.stat().st_size, - "sha256": sha256_file(args.surgery_manifest), - }, "components": { "text": component_record(args.text, text_validation), "mmproj": component_record(args.mmproj, mmproj_validation), From fb38841f6ca433f835f72963357f477c0c8388f4 Mon Sep 17 00:00:00 2001 From: JJJYmmm <1650675829@qq.com> Date: Wed, 2 Sep 2026 03:24:30 +0800 Subject: [PATCH 10/11] docs: add StarVLA setup and model zoo --- README.md | 95 ++++++++++++++++++++++++++++++++++----- README_ZH.md | 82 +++++++++++++++++++++++++++++++-- robot_server/README.md | 65 +++++++++++++++++++++++++-- robot_server/README_ZH.md | 77 ++++++++++++++++++++++++++----- 4 files changed, 291 insertions(+), 28 deletions(-) diff --git a/README.md b/README.md index 055e826..29495dd 100644 --- a/README.md +++ b/README.md @@ -54,8 +54,13 @@ We also provide two tools to support robot model development: git clone https://github.com/Robot-cpp/robot.cpp cd robot.cpp git submodule update --init --recursive +./tools/apply_patches.sh ``` +The launch scripts below configure and build `model-server` automatically. For +a manual StarVLA build, enable `ROBOT_CPP_BUILD_STARVLA`; see the +[Robot Server build instructions](robot_server/README.md#manual-build). + This section introduces three usage paths to help you quickly understand the repository: * Starting `model-server` and connecting it to a minimal dummy `model-client`. @@ -95,6 +100,14 @@ After downloading, run `model-server` like this: For general local setups, we provide ready-to-use build-and-launch shells for three platforms. You can modify the environment variables inside the scripts, or override them directly with `export`. See [robot_server/README.md](robot_server/README.md) for details. +For example, from the repository root on Linux with CUDA: + +```bash +export ROBOT_CPP_ROOT="$PWD" +export GGUF_DIR=/path/to/smolvla-so101-fp32 +bash robot_server/shell/launch_robot_server_linux_cuda.sh +``` + | Backend | macOS | Linux | Windows | | ------- | ------------------------------------------------------- | -------------------------------------------------------- | ----------------------------------------------------------- | | CUDA | - | `robot_server/shell/launch_robot_server_linux_cuda.sh` | `robot_server/shell/launch_robot_server_windows_cuda.bat` | @@ -127,7 +140,7 @@ We provide a build-to-run example in `robot_client/shell/cpp_client_example.sh`. | `ROBOT_CPP_ROOT` | unset; required | Repository root. | | `BUILD_DIR` | `${ROBOT_CPP_ROOT}/build_robot_client` | C++ client CMake build directory. | | `PORT` | `5555` | Server port used by the client. | -| `BUILD_CLIENT` | `0` | Whether to force rebuild the client. Set to`1` to rebuild even if the binary already exists. | +| `BUILD_CLIENT` | `0` | Whether to force rebuild the client. Set to `1` to rebuild even if the binary already exists. | | `CMAKE_BIN` | `cmake` | CMake command path, useful for selecting a custom CMake binary. | Then run: @@ -150,20 +163,33 @@ See the [SO-101 deployment guide](eval/lerobot_so101/README.md). ## ⚡ Performance -We benchmark Robot.cpp on several platforms. Each measurement uses 5 warmup runs and 100 loop runs. The reported latency is the average time from receiving the image, through preprocessing and forward inference, to producing a usable action chunk, measured in milliseconds. All state projectors remain in f32 precision. +We benchmark Robot.cpp on several platforms. Each measurement uses 5 warmup runs and 100 loop runs. The reported latency is the average time from receiving the image, through preprocessing and forward inference, to producing a usable action chunk, measured in milliseconds. State projectors, where present, remain in f32 precision. For the LIBERO setting, the input contains two 256x256 images and an 8-dimensional state. For the SO-101 real-robot setting, the input contains one 224x224 image and a 6-dimensional state. For SmolVLA preprocessing, we follow the official default setting: images are first resized to 512x512. -| Model | Mac M4 Pro (CPU) | Mac M4 Pro (Metal) | RTX 4090 | RTX 3060 | A100 | Jetson AGX Orin | -| ---------------------- | ---------------: | -----------------: | -------: | ----------: | ---: | --------------: | -| smolvla@libero (bf16*) | 527 | 216 | 28 | 116 | 43 | 282 | -| smolvla@libero (f32) | 577 | 236 | 32 | 142 | 42 | 299 | -| smolvla@so-101 (bf16*) | 339 | 145 | 23 | 77 | 36 | 184 | -| smolvla@so-101 (f32) | 396 | 158 | 24 | 92 | 34 | 200 | -| pi0@libero (f32) | 1839 | 710 | 83 | OOM/offload | 71 | 956 | -| pi0@libero (bf16*) | 1954 | 635 | 57 | 267 | 66 | 498 | +For StarVLA, the input contains one 224x224 image and no robot state. Qwen and +the multimodal projector use bf16; OFT, GR00T, PI, and PI_v3 policies use f32. +FAST stores its action codec in the policy GGUF. +The StarVLA A100 results use an A100-PCIE-40GB with 8 CPU threads, +`n_ctx=2048`, `n_batch=2048`, and noise seed 0. + +| Model | Mac M4 Pro (CPU) | Mac M4 Pro (Metal) | RTX 4090 | RTX 3060 | A100 | Jetson AGX Orin | +| ----------------------------- | ---------------: | -----------------: | -------: | ----------: | ---: | --------------: | +| smolvla@libero (bf16*) | 527 | 216 | 28 | 116 | 43 | 282 | +| smolvla@libero (f32) | 577 | 236 | 32 | 142 | 42 | 299 | +| smolvla@so-101 (bf16*) | 339 | 145 | 23 | 77 | 36 | 184 | +| smolvla@so-101 (f32) | 396 | 158 | 24 | 92 | 34 | 200 | +| pi0@libero (f32) | 1839 | 710 | 83 | OOM/offload | 71 | 956 | +| pi0@libero (bf16*) | 1954 | 635 | 57 | 267 | 66 | 498 | +| starvla/oft@bridge | - | - | - | - | 50 | - | +| starvla/groot@bridge | - | - | - | - | 54 | - | +| starvla/pi_v3@bridge | - | - | - | - | 112 | - | +| starvla/qwen25_oft@bridge | - | - | - | - | 42 | - | +| starvla/qwen25_groot@bridge | - | - | - | - | 51 | - | +| starvla/qwen25_pi@bridge | - | - | - | - | 101 | - | +| starvla/qwen25_fast@bridge | - | - | - | - | 386 | - | > `bf16*`: on Mac, f16 results are used in place of bf16 because current Mac bf16 support is not ideal. > `OOM/offload`: pi0@libero (f32) runs out of memory on RTX 3060 and triggers offload, so we do not report a latency number for now. @@ -230,6 +256,55 @@ This section lists converted GGUF models that can be used directly with `model-s f32 pi0-libero-f32 + + StarVLA Qwen3-VL OFT + Bridge + StarVLA/Qwen3VL-OFT-Bridge-RT-1 + bf16 + f32 policy + starvla-qwen3-oft-bridge-bf16 + + + StarVLA Qwen3-VL GR00T + Bridge + StarVLA/Qwen3VL-GR00T-Bridge-RT-1 + bf16 + f32 policy + starvla-qwen3-groot-bridge-bf16 + + + StarVLA Qwen3-VL PI_v3 + Bridge + StarVLA/Qwen3VL-PI_v3-Bridge-RT_1 + bf16 + f32 policy + starvla-qwen3-pi-v3-bridge-bf16 + + + StarVLA Qwen2.5-VL OFT + Bridge + StarVLA/Qwen-OFT-Bridge-RT-1 + bf16 + f32 policy + starvla-qwen25-oft-bridge-bf16 + + + StarVLA Qwen2.5-VL GR00T + Bridge + StarVLA/Qwen-GR00T-Bridge-RT-1 + bf16 + f32 policy + starvla-qwen25-groot-bridge-bf16 + + + StarVLA Qwen2.5-VL PI + Bridge + StarVLA/Qwen-PI-Bridge-RT-1 + bf16 + f32 policy + starvla-qwen25-pi-bridge-bf16 + + + StarVLA Qwen2.5-VL FAST + Bridge + StarVLA/Qwen-FAST-Bridge-RT-1 + bf16 + codec + starvla-qwen25-fast-bridge-bf16 + diff --git a/README_ZH.md b/README_ZH.md index 78c5e4a..edbe945 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -54,8 +54,13 @@ Robot.cpp是一个轻量化的on-device机器人模型推理框架，在llama.cp git clone https://github.com/Robot-cpp/robot.cpp cd robot.cpp git submodule update --init --recursive +./tools/apply_patches.sh ``` +下文的启动脚本会自动配置并编译 `model-server`。手动构建 StarVLA 时需要开启 +`ROBOT_CPP_BUILD_STARVLA`，详见 +[Robot Server 构建说明](robot_server/README_ZH.md#手动构建)。 + 我们介绍三类使用案例来帮助你快速了解本仓库： * model-server的启动，其与最小dummy model-client通信的案例。 @@ -95,6 +100,14 @@ git submodule update --init --recursive 对于更加一般的情况，我们也提供了三个平台的开箱即用编译+启动的shell，可以通过修改shell里的环境变量，或者直接export的形式来快速在本机实现启动。详情参见 [robot_server/README_ZH.md](robot_server/README_ZH.md) +例如，在 Linux CUDA 环境中从仓库根目录运行： + +```bash +export ROBOT_CPP_ROOT="$PWD" +export GGUF_DIR=/path/to/smolvla-so101-fp32 +bash robot_server/shell/launch_robot_server_linux_cuda.sh +``` + | Backend | macOS | Linux | Windows | | ------- | ------------------------------------------------------- | -------------------------------------------------------- | ----------------------------------------------------------- | | CUDA | - | `robot_server/shell/launch_robot_server_linux_cuda.sh` | `robot_server/shell/launch_robot_server_windows_cuda.bat` | @@ -127,7 +140,7 @@ python robot_client/examples/python/minimal_example.py | `ROBOT_CPP_ROOT` | 无，必须设置 | 仓库根目录。 | | `BUILD_DIR` | `${ROBOT_CPP_ROOT}/build_robot_client` | C++ client 的 CMake build 目录 | | `PORT` | `5555` | client 连接的 server port | -| `BUILD_CLIENT` | `0` | 是否强制重新build client。设为`1` 时即使 binary 已存在也会重新 build | +| `BUILD_CLIENT` | `0` | 是否强制重新build client。设为 `1` 时即使 binary 已存在也会重新 build | | `CMAKE_BIN` | `cmake` | 使用的 CMake 命令路径，可用于指定自定义 CMake | 然后运行下面的bash： @@ -156,6 +169,11 @@ bash robot_client/shell/cpp_client_example.sh 其中对于smolvla的preprocess设定，参考官方的基本设定，即首先会将图片变成512*512。 +StarVLA 使用一张 224x224 图像且不输入 robot state。Qwen 和 multimodal projector +使用 bf16，OFT、GR00T、PI 和 PI_v3 policy 使用 f32；FAST 的 policy GGUF 保存 action +codec。A100 数据在 A100-PCIE-40GB、8 个 CPU 线程、`n_ctx=2048`、`n_batch=2048` 和 +noise seed 0 下测得。 + | Model | Mac M4 Pro (CPU) | Mac M4 Pro (Metal) | RTX 4090 | RTX 3060 | A100 | Jetson AGX Orin | | ---------------------- | ---------------: | -----------------: | -------: | ----------: | ---: | --------------: | | smolvla@libero (bf16*) | 527 | 216 | 28 | 116 | 43 | 282 | @@ -164,15 +182,24 @@ bash robot_client/shell/cpp_client_example.sh | smolvla@so-101 (f32) | 396 | 158 | 24 | 92 | 34 | 200 | | pi0@libero (f32) | 1839 | 710 | 83 | OOM/offload | 71 | 956 | | pi0@libero (bf16*) | 1954 | 635 | 57 | 267 | 66 | 498 | +| starvla/oft@bridge | - | - | - | - | 50 | - | +| starvla/groot@bridge | - | - | - | - | 54 | - | +| starvla/pi_v3@bridge | - | - | - | - | 112 | - | +| starvla/qwen25_oft@bridge | - | - | - | - | 42 | - | +| starvla/qwen25_groot@bridge | - | - | - | - | 51 | - | +| starvla/qwen25_pi@bridge | - | - | - | - | 101 | - | +| starvla/qwen25_fast@bridge | - | - | - | - | 386 | - | > `bf16*`：在 Mac上使用 f16 结果替代 bf16，因为当前 Mac对 bf16 的支持不够好。 > `OOM/offload`：pi0@libero (f32) 在 RTX 3060 上会 OOM 并触发 offload，因此暂时不报告 latency 数值。 --- -## 🧩 model-zoo +## 🧩 Model Zoo -这里整理一些已经转换好的 GGUF 模型，可以直接配合 `model-server` 做smoke test，以方便quick start！但针对自己的实际场景，我们推荐使用[hf2gguf](tools/hf2gguf/README_ZH.md)来生成自己的GGUF model！并且对于不同的部分，您还可以自定义不同的精度，来实现不同部分的精度组合（事实上，不同部分的最优精度通常是不同的），我们的例子中，state proj始终保持f32精度，其他的gguf随着precision精度变化而变化，您可以自行组合，探索更好更高效的性能tradeoff！ +下表列出可直接配合 `model-server` 使用的 GGUF 模型。实际部署时，建议使用 +[`hf2gguf`](tools/hf2gguf/README_ZH.md) 转换自己的 checkpoint。各组件可以分别选择 +精度；表中示例的 state projector 固定为 f32，其余组件采用标注的精度。 @@ -230,6 +257,55 @@ bash robot_client/shell/cpp_client_example.sh + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
f32 pi0-libero-f32
StarVLA Qwen3-VL OFTBridgeStarVLA/Qwen3VL-OFT-Bridge-RT-1bf16 + f32 policystarvla-qwen3-oft-bridge-bf16
StarVLA Qwen3-VL GR00TBridgeStarVLA/Qwen3VL-GR00T-Bridge-RT-1bf16 + f32 policystarvla-qwen3-groot-bridge-bf16
StarVLA Qwen3-VL PI_v3BridgeStarVLA/Qwen3VL-PI_v3-Bridge-RT_1bf16 + f32 policystarvla-qwen3-pi-v3-bridge-bf16
StarVLA Qwen2.5-VL OFTBridgeStarVLA/Qwen-OFT-Bridge-RT-1bf16 + f32 policystarvla-qwen25-oft-bridge-bf16
StarVLA Qwen2.5-VL GR00TBridgeStarVLA/Qwen-GR00T-Bridge-RT-1bf16 + f32 policystarvla-qwen25-groot-bridge-bf16
StarVLA Qwen2.5-VL PIBridgeStarVLA/Qwen-PI-Bridge-RT-1bf16 + f32 policystarvla-qwen25-pi-bridge-bf16
StarVLA Qwen2.5-VL FASTBridgeStarVLA/Qwen-FAST-Bridge-RT-1bf16 + codecstarvla-qwen25-fast-bridge-bf16
diff --git a/robot_server/README.md b/robot_server/README.md index 7b43f6a..37f01fb 100644 --- a/robot_server/README.md +++ b/robot_server/README.md @@ -9,6 +9,17 @@ the Python/C++ clients. The server keeps one robot policy model loaded in-process and returns an action chunk for each prediction request. There are two main ways to use it. +Before building from source, initialize the submodules and apply the repository +patches from the project root: + +```bash +git submodule update --init --recursive +./tools/apply_patches.sh +``` + +Source builds require CMake 3.16 or newer and a C/C++ compiler. CUDA builds also +require the CUDA Toolkit and a working `nvcc`. + ## Method 1: One-Command Build and Run We provide one-command build-and-run scripts for several mainstream platforms @@ -21,6 +32,9 @@ under `robot_server/shell`. After running a script, the corresponding | CPU | `robot_server/shell/launch_robot_server_mac_cpu.sh` | `robot_server/shell/launch_robot_server_linux_cpu.sh` | `robot_server/shell/launch_robot_server_windows_cpu.bat` | | Metal | `robot_server/shell/launch_robot_server_mac_metal.sh` | - | - | +The Linux CUDA script supports SmolVLA, pi0, and StarVLA. The other launch +scripts currently support SmolVLA and pi0. + ### Set Variables Before running a script, configure the variables below as needed. @@ -30,14 +44,18 @@ Common variables: | Variable | Description | Default | | ------------------ | ----------------------------------------------------------------------------------- | ----------------------------------------------------------------------------- | | `ROBOT_CPP_ROOT` | Repository root. | Must be set explicitly | -| `MODEL_TYPE` | Model type,`smolvla` or `pi0`. | `smolvla` | +| `MODEL_TYPE` | Model type: `smolvla`, `pi0`, or `starvla`. | `smolvla` | | `GGUF_DIR` | Directory containing GGUF files. | Must be set explicitly | -| `BUILD_DIR` | CMake build directory. | macOS / Linux defaults are organized as`build_{mac/linux}_{cpu/metal/cuda}` | +| `BUILD_DIR` | CMake build directory. | macOS / Linux defaults use `build_{mac/linux}_{cpu/metal/cuda}` | +| `HOST` | Server listen address. | `127.0.0.1` | | `PORT` | Server listen port. | `5555` | | `THREADS` | Inference thread count. | `8` | -| `TASK` | Language input describing the task. | `grab the block.` | +| `N_BATCH` | LLM batch size. | `512` | +| `N_CTX` | LLM context size. | `2048` | +| `TASK` | Compatibility fallback; each client request supplies the actual task. | `grab the block.` | | `NOISE_SEED` | Action noise seed. | `-1` | -| `SKIP_BUILD` | Whether to skip configure/build. Set to`1` to directly launch an existing binary. | `0` | +| `VERBOSITY` | Model log verbosity. | `0` | +| `SKIP_BUILD` | Whether to skip configure/build. Set to `1` to launch an existing binary. | `0` | | `CMAKE_BIN` | CMake executable. | `cmake` | SmolVLA variables: @@ -61,6 +79,14 @@ pi0 variables: | `STATE_GGUF` | Full path to the pi0 state GGUF. | `${GGUF_DIR}/${MODEL_BASENAME}.state.gguf` | | `ACTION_DECODER_GGUF` | Full path to the pi0 action decoder GGUF. | `${GGUF_DIR}/${MODEL_BASENAME}.action_decoder.gguf` | +StarVLA variables (Linux CUDA): + +| Variable | Description | Default | +| ------------- | ------------------------------------- | -------- | +| `LLM_GGUF` | Full path to the Qwen text GGUF. | Required | +| `MMPROJ_GGUF` | Full path to the Qwen vision GGUF. | Required | +| `POLICY_GGUF` | Full path to the StarVLA policy GGUF. | Required | + ### Invocation Run the macOS / Linux `.sh` scripts with `bash`: @@ -74,6 +100,37 @@ bash robot_server/shell/launch_robot_server_linux_cuda.sh Windows uses the `.bat` scripts. +For example, run the Qwen3-VL OFT bundle from the Model Zoo on Linux CUDA: + +```bash +export ROBOT_CPP_ROOT="$PWD" +export GGUF_DIR=/path/to/starvla-qwen3-oft-bridge-bf16 +export MODEL_TYPE=starvla +export LLM_GGUF="${GGUF_DIR}/qwen-oft-bf16.gguf" +export MMPROJ_GGUF="${GGUF_DIR}/mmproj-oft-bf16.gguf" +export POLICY_GGUF="${GGUF_DIR}/starvla-oft-policy-fp32.gguf" +bash robot_server/shell/launch_robot_server_linux_cuda.sh +``` + +The script automatically configures CMake with +`ROBOT_CPP_BUILD_STARVLA=ON` when `MODEL_TYPE=starvla`. + +### Manual Build + +To build a StarVLA-enabled server without a launch script: + +```bash +cmake -S . -B build_cuda \ + -DCMAKE_BUILD_TYPE=Release \ + -DGGML_CUDA=ON \ + -DROBOT_CPP_BUILD_ROBOT_SERVER=ON \ + -DROBOT_CPP_BUILD_STARVLA=ON +cmake --build build_cuda --target model-server -j +``` + +`ROBOT_CPP_BUILD_STARVLA` defaults to `OFF`; SmolVLA and pi0 builds do not need +it. + ### Troubleshooting - `Tell CMake where to find the compiler by setting either the environment variable "CUDACXX" or the CMake cache entry CMAKE_CUDA_COMPILER to the full path to the compiler, or to the compiler name if it is in the PATH.` diff --git a/robot_server/README_ZH.md b/robot_server/README_ZH.md index 61a708b..3d6f2fa 100644 --- a/robot_server/README_ZH.md +++ b/robot_server/README_ZH.md @@ -4,11 +4,22 @@ # Robot Server -`robot_server` 提供了 `model-server` 与 Python/C++ 客户端使用的轻量级 TCP 协议。服务端会在进程内常驻加载一个机器人策略模型，并针对每个预测请求返回一段 action chunk。具体有两种使用方式 +`robot_server` 提供了 `model-server` 与 Python/C++ 客户端使用的轻量级 TCP 协议。服务端会在进程内常驻加载一个机器人策略模型，并针对每个预测请求返回一段 action chunk。具体有两种使用方式。 -## 方法1：一键编译与运行 +从源码构建前，先在仓库根目录初始化子模块并应用项目补丁： -我们在`robot_server/shell`下提供了几种主流平台的一键编译运行的脚本，运行之后，将直接开始运行对应的`model-server`开始监听 +```bash +git submodule update --init --recursive +./tools/apply_patches.sh +``` + +源码构建需要 CMake 3.16 或更高版本及 C/C++ 编译器。CUDA 构建还需要 CUDA Toolkit +和可用的 `nvcc`。 + +## 方法 1：一键编译与运行 + +`robot_server/shell` 提供多个平台的编译启动脚本。脚本会配置、编译并启动对应的 +`model-server`。 | Backend | macOS | Linux | Windows | | ------- | ------------------------------------------------------- | -------------------------------------------------------- | ----------------------------------------------------------- | @@ -16,23 +27,29 @@ | CPU | `robot_server/shell/launch_robot_server_mac_cpu.sh` | `robot_server/shell/launch_robot_server_linux_cpu.sh` | `robot_server/shell/launch_robot_server_windows_cpu.bat` | | Metal | `robot_server/shell/launch_robot_server_mac_metal.sh` | - | - | +Linux CUDA 脚本支持 SmolVLA、pi0 和 StarVLA；其他启动脚本当前支持 SmolVLA 和 pi0。 + ### 设置变量 -在运行之前，我们需要在脚本中设置一些设定上的变量，具体而言有以下变量可以按需设置。 +运行脚本前，按需设置以下环境变量。 公共变量： | 变量 | 说明 | 默认值 | | ------------------ | --------------------------------------------------------- | --------------------------------------------------------------- | | `ROBOT_CPP_ROOT` | 仓库根目录 | 必须显式设置 | -| `MODEL_TYPE` | 模型类型，可选`smolvla` / `pi0` | `smolvla` | +| `MODEL_TYPE` | 模型类型，可选 `smolvla`、`pi0` 或 `starvla` | `smolvla` | | `GGUF_DIR` | GGUF 文件所在目录 | 必须显式设置 | -| `BUILD_DIR` | CMake build 目录 | macOS / Linux 默认按`build_{mac/linux}_{cpu/metal/cuda}` 组织 | +| `BUILD_DIR` | CMake build 目录 | macOS / Linux 默认使用 `build_{mac/linux}_{cpu/metal/cuda}` | +| `HOST` | server 监听地址 | `127.0.0.1` | | `PORT` | server 监听端口 | `5555` | | `THREADS` | 推理线程数 | `8` | -| `TASK` | 语言输入，描述任务 | `grab the block.` | +| `N_BATCH` | LLM batch size | `512` | +| `N_CTX` | LLM context size | `2048` | +| `TASK` | 兼容用默认值；实际 task 由每个 client 请求提供 | `grab the block.` | | `NOISE_SEED` | action noise seed | `-1` | -| `SKIP_BUILD` | 是否跳过 configure/build，设为`1` 时直接启动已有 binary | `0` | +| `VERBOSITY` | 模型日志级别 | `0` | +| `SKIP_BUILD` | 是否跳过 configure/build；设为 `1` 时直接启动已有 binary | `0` | | `CMAKE_BIN` | CMake 可执行文件 | `cmake` | SmolVLA 变量： @@ -56,6 +73,14 @@ pi0 变量： | `STATE_GGUF` | pi0 state GGUF 完整路径 | `${GGUF_DIR}/${MODEL_BASENAME}.state.gguf` | | `ACTION_DECODER_GGUF` | pi0 action decoder GGUF 完整路径 | `${GGUF_DIR}/${MODEL_BASENAME}.action_decoder.gguf` | +StarVLA 变量（Linux CUDA）： + +| 变量 | 说明 | 默认值 | +| ------------- | ---------------------------- | -------- | +| `LLM_GGUF` | Qwen text GGUF 完整路径 | 必须设置 | +| `MMPROJ_GGUF` | Qwen vision GGUF 完整路径 | 必须设置 | +| `POLICY_GGUF` | StarVLA policy GGUF 完整路径 | 必须设置 | + ### 调用方式 macOS / Linux 的 `.sh` 脚本直接用 `bash` 运行。 @@ -67,7 +92,37 @@ bash robot_server/shell/launch_robot_server_linux_cpu.sh bash robot_server/shell/launch_robot_server_linux_cuda.sh ``` -Windows 的 `.bat` 脚本： +Windows 直接运行对应的 `.bat` 脚本。 + +例如，在 Linux CUDA 上运行 Model Zoo 中的 Qwen3-VL OFT bundle： + +```bash +export ROBOT_CPP_ROOT="$PWD" +export GGUF_DIR=/path/to/starvla-qwen3-oft-bridge-bf16 +export MODEL_TYPE=starvla +export LLM_GGUF="${GGUF_DIR}/qwen-oft-bf16.gguf" +export MMPROJ_GGUF="${GGUF_DIR}/mmproj-oft-bf16.gguf" +export POLICY_GGUF="${GGUF_DIR}/starvla-oft-policy-fp32.gguf" +bash robot_server/shell/launch_robot_server_linux_cuda.sh +``` + +当 `MODEL_TYPE=starvla` 时，脚本会自动使用 +`ROBOT_CPP_BUILD_STARVLA=ON` 配置 CMake。 + +### 手动构建 + +不使用启动脚本时，可按以下方式构建支持 StarVLA 的 server： + +```bash +cmake -S . -B build_cuda \ + -DCMAKE_BUILD_TYPE=Release \ + -DGGML_CUDA=ON \ + -DROBOT_CPP_BUILD_ROBOT_SERVER=ON \ + -DROBOT_CPP_BUILD_STARVLA=ON +cmake --build build_cuda --target model-server -j +``` + +`ROBOT_CPP_BUILD_STARVLA` 默认为 `OFF`；构建 SmolVLA 和 pi0 时不需要开启。 ### 故障排查 @@ -80,9 +135,9 @@ export CUDACXX=/usr/local/cuda-12.4/bin/nvcc export PATH=/usr/local/cuda-12.4/bin:$PATH ``` -## 方法2：直接下载预编译发布 +## 方法 2：直接下载预编译发布 -从release page下载之后，运行以下命令 +从 release page 下载后，运行以下命令。 ### 启动 SmolVLA From 2cea6a7e25992753a1f0180b1f69be36edc67105 Mon Sep 17 00:00:00 2001 From: JJJYmmm <1650675829@qq.com> Date: Wed, 2 Sep 2026 03:25:20 +0800 Subject: [PATCH 11/11] eval: add SimplerEnv StarVLA evaluation --- README.md | 6 +- README_ZH.md | 6 +- eval/README.md | 6 + eval/README_ZH.md | 4 + eval/simpler_env/README.md | 97 +++++ eval/simpler_env/README_ZH.md | 97 +++++ eval/simpler_env/__init__.py | 1 + eval/simpler_env/environment.yaml | 19 + eval/simpler_env/policy/__init__.py | 1 + eval/simpler_env/policy/model_server.py | 234 ++++++++++ eval/simpler_env/runners/__init__.py | 1 + eval/simpler_env/runners/latency_starvla.py | 424 +++++++++++++++++++ eval/simpler_env/runners/run_model_server.py | 404 ++++++++++++++++++ eval/simpler_env/scripts/run_model_server.sh | 81 ++++ eval/simpler_env/utils/__init__.py | 1 + eval/simpler_env/utils/environment.py | 230 ++++++++++ 16 files changed, 1610 insertions(+), 2 deletions(-) create mode 100644 eval/simpler_env/README.md create mode 100644 eval/simpler_env/README_ZH.md create mode 100644 eval/simpler_env/__init__.py create mode 100644 eval/simpler_env/environment.yaml create mode 100644 eval/simpler_env/policy/__init__.py create mode 100644 eval/simpler_env/policy/model_server.py create mode 100644 eval/simpler_env/runners/__init__.py create mode 100644 eval/simpler_env/runners/latency_starvla.py create mode 100755 eval/simpler_env/runners/run_model_server.py create mode 100755 eval/simpler_env/scripts/run_model_server.sh create mode 100644 eval/simpler_env/utils/__init__.py create mode 100644 eval/simpler_env/utils/environment.py diff --git a/README.md b/README.md index 29495dd..4e02885 100644 --- a/README.md +++ b/README.md @@ -151,7 +151,9 @@ bash robot_client/shell/cpp_client_example.sh ### 🧪 Using model-server in simulation, using LIBERO as the example -See the [LIBERO simulation evaluation guide](eval/libero/README.md). +See the [LIBERO simulation evaluation guide](eval/libero/README.md). To run +local StarVLA GGUF models on the WidowX Bridge tasks, see the +[SimplerEnv Bridge guide](eval/simpler_env/README.md). ### 🦾 Using model-server on real hardware, using SO-101 as the example @@ -344,6 +346,7 @@ robot.cpp/ ├── eval/ │ ├── base_platform.py # Shared base class for real-robot platforms │ ├── libero/ # LIBERO simulation evaluation +│ ├── simpler_env/ # SimplerEnv WidowX / Bridge evaluation │ └── lerobot_so101/ # SO-101 real-robot scripts and examples └── third_party/ ├── llama.cpp/ # ggml / llama.cpp backend @@ -400,6 +403,7 @@ Robot.cpp's design and implementation benefit from several excellent open-source * [llama.cpp](https://github.com/ggerganov/llama.cpp): provides lightweight local inference, the GGML/GGUF ecosystem, and cross-platform backend foundations. This project continues building robot model inference capabilities on top of its engineering philosophy and low-level runtime. * [LeRobot](https://github.com/huggingface/lerobot): provides reference implementations for robot data, policy training, and real-robot integration. The SO-101 real-robot example and parts of the evaluation flow in this project are inspired by the LeRobot ecosystem. * [LIBERO](https://github.com/Lifelong-Robot-Learning/LIBERO): provides robot simulation tasks and evaluation benchmarks. The LIBERO simulation evaluation flow in this project is based on its task environments and benchmark design. +* [SimplerEnv](https://github.com/simpler-env/SimplerEnv): provides real-to-sim robot evaluation environments. The StarVLA Bridge success-rate evaluation uses its WidowX task suite and official visual-matching assets. * [OpenPI](https://github.com/Physical-Intelligence/openpi): provides the pi0 policy model and related open-source implementation. The pi0 runtime, conversion, and evaluation work in this project references OpenPI's model design. Thanks to these projects and communities for their contributions to robot learning and on-device inference. diff --git a/README_ZH.md b/README_ZH.md index edbe945..0d61dfa 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -151,7 +151,9 @@ bash robot_client/shell/cpp_client_example.sh ### 🧪 model-server在仿真平台上的使用（以LIBERO为例） -详见 [LIBERO 仿真评测说明](eval/libero/README_ZH.md)。 +详见 [LIBERO 仿真评测说明](eval/libero/README_ZH.md)。在 WidowX Bridge 任务上运行 +StarVLA，并比较 Python checkpoint 与 GGUF 的方法见 +[SimplerEnv Bridge 说明](eval/simpler_env/README_ZH.md)。 ### 🦾 model-server在真机平台上的使用（以SO-101为例） @@ -345,6 +347,7 @@ robot.cpp/ ├── eval/ │ ├── base_platform.py # 真机 platform 的统一基类 │ ├── libero/ # LIBERO 仿真评测 +│ ├── simpler_env/ # SimplerEnv WidowX / Bridge 仿真评测 │ └── lerobot_so101/ # SO-101 真机相关脚本与示例 └── third_party/ ├── llama.cpp/ # ggml / llama.cpp 后端 @@ -401,6 +404,7 @@ robot.cpp 的设计与实现受益于多个优秀的开源项目： * [llama.cpp](https://github.com/ggerganov/llama.cpp)：提供了轻量化本地推理、GGML/GGUF 生态与跨平台后端基础，本项目在其工程哲学和底层能力上继续构建机器人模型推理框架。 * [LeRobot](https://github.com/huggingface/lerobot)：提供了机器人数据、策略训练与真实机器人接入的参考实现，本项目的 SO-101 真机示例与部分评测流程参考了 LeRobot 生态。 * [LIBERO](https://github.com/Lifelong-Robot-Learning/LIBERO)：提供了机器人仿真任务与评测基准，本项目的 LIBERO 仿真评测流程基于其任务环境与 benchmark 设计。 +* [SimplerEnv](https://github.com/simpler-env/SimplerEnv)：提供了 real-to-sim 机器人评测环境，本项目的 StarVLA Bridge 成功率评测使用其 WidowX 任务集与官方 visual-matching 资产。 * [OpenPI](https://github.com/Physical-Intelligence/openpi)：提供了pi0策略模型与相关开源实现，本项目的 pi0 相关 runtime、转换与评测工作参考了 OpenPI 的模型设计。 感谢这些项目和社区为机器人学习与端侧推理生态做出的贡献。 diff --git a/eval/README.md b/eval/README.md index 9cf9e75..0d9e672 100644 --- a/eval/README.md +++ b/eval/README.md @@ -22,6 +22,7 @@ The repo root [README.md](../README.md) describes the three-layer layout: eval/ ├── base_platform.py # Shared base class for real-robot / sim platforms ├── libero/ # LIBERO sim benchmark (multi-camera, batch rollout) +├── simpler_env/ # SimplerEnv WidowX / Bridge closed-loop benchmark └── lerobot_so101/ # SO-101 real-robot sync closed-loop example ``` @@ -30,12 +31,16 @@ eval/ | Directory | Scenario | Notes | |---|---|---| | [`libero/`](libero/README.md) | Sim eval | LIBERO benchmark with C++ policy rollout and LeRobot baseline. [中文](libero/README_ZH.md) | +| [`simpler_env/`](simpler_env/README.md) | Sim eval | Runs StarVLA Python and GGUF on the SimplerEnv WidowX Bridge tasks. [中文](simpler_env/README_ZH.md) | | [`lerobot_so101/`](lerobot_so101/README.md) | Real robot | SO-101 follower + single-camera observe → predict → act loop. [中文](lerobot_so101/README_ZH.md) | The two examples are organized slightly differently: - **SO-101** follows the standard `BasePlatform` + `RobotPolicy` + `SyncControlLoop` path—use it as the template for new real-robot platforms. - **LIBERO** implements a dedicated observation adapter under `eval/libero/policy/` (multi-camera, state packing, sim rollout) and does **not** inherit `BasePlatform`—use it as a reference for new **sim benchmarks**. +- **SimplerEnv** follows the same dedicated runner pattern and implements the + WidowX action transform, temporal ensemble, normalization profile, and Bridge + task settings for both Python and C++ runs. ## Standard closed-loop data flow @@ -237,6 +242,7 @@ LIBERO’s [`ModelServerPolicy`](libero/policy/model_server.py) is an example of - [SO-101 real-robot guide](lerobot_so101/README.md) · [中文](lerobot_so101/README_ZH.md) - [LIBERO sim eval](libero/README.md) · [中文](libero/README_ZH.md) +- [SimplerEnv Bridge eval](simpler_env/README.md) · [中文](simpler_env/README_ZH.md) - [robot_server launch and protocol](../robot_server/README.md) - [robot_client and policy](../robot_client/README.md) - [Adding a new model runtime](../src/README.md) diff --git a/eval/README_ZH.md b/eval/README_ZH.md index eb1e781..99784b1 100644 --- a/eval/README_ZH.md +++ b/eval/README_ZH.md @@ -24,6 +24,7 @@ eval/ ├── base_platform.py # 真机 / 仿真 platform 的统一基类 ├── libero/ # LIBERO 仿真 benchmark（多相机、批量 rollout） +├── simpler_env/ # SimplerEnv WidowX / Bridge 闭环 benchmark └── lerobot_so101/ # SO-101 真机同步闭环示例 ``` @@ -33,6 +34,7 @@ eval/ | 目录 | 场景 | 说明 | | ---------------------------------------------- | ---- | ------------------------------------------------------------------------------------------ | | `[libero/](libero/README_ZH.md)` | 仿真评测 | 面向 LIBERO benchmark，含 C++ policy rollout 与 LeRobot baseline 对比。[English](libero/README.md) | +| `[simpler_env/](simpler_env/README_ZH.md)` | 仿真评测 | 在 SimplerEnv WidowX Bridge 任务上运行并比较 StarVLA Python 与 GGUF。[English](simpler_env/README.md) | | `[lerobot_so101/](lerobot_so101/README_ZH.md)` | 真机闭环 | SO-101 follower + 单相机的 observe → predict → act 同步控制。[English](lerobot_so101/README.md) | @@ -40,6 +42,7 @@ eval/ - **SO-101** 走标准 `BasePlatform` + `RobotPolicy` + `SyncControlLoop` 路径，适合作为新增真机 platform 的模板。 - **LIBERO** 在 `eval/libero/policy/` 里实现了专用的 observation 适配（多相机、state 拼接、仿真 rollout），不继承 `BasePlatform`，适合作为新增 **仿真 benchmark** 的参考。 +- **SimplerEnv** 沿用专用 runner 结构，为 Python 和 C++ 实现相同的 WidowX action 变换、时序集成、normalization profile 与 Bridge 任务设置。 ## 标准闭环数据流 @@ -244,6 +247,7 @@ LIBERO 的 `[ModelServerPolicy](libero/policy/model_server.py)` 即为自定义 - [SO-101 真机使用说明](lerobot_so101/README_ZH.md) - [LIBERO 仿真评测说明](libero/README_ZH.md) +- [SimplerEnv Bridge 仿真评测说明](simpler_env/README_ZH.md) - [robot_server 启动与协议](../robot_server/README_ZH.md) - [robot_client 与 policy](../robot_client/README.md) - [新增模型 runtime](../src/README_ZH.md) diff --git a/eval/simpler_env/README.md b/eval/simpler_env/README.md new file mode 100644 index 0000000..b22dc91 --- /dev/null +++ b/eval/simpler_env/README.md @@ -0,0 +1,97 @@ +# SimplerEnv WidowX Bridge Eval + +This directory evaluates the robot.cpp StarVLA GGUF runtime on the SimplerEnv +WidowX Bridge tasks. + +## Protocol + +- Four Bridge tasks with object episodes `0..23` +- At most 120 steps per episode at 5 Hz +- Visual-matching RGB overlay resized to 224x224 with OpenCV `INTER_AREA` +- One action chunk per step with the official seven-prediction adaptive ensemble + +A full run contains 96 rollouts. The result reports overall and per-task success +rates; subset runs use `partial` coverage. + +## Setup + +Convert a checkpoint as described in the +[StarVLA guide](../../tools/hf2gguf/starvla/README.md) and build the CUDA runtime. + +The environment uses these revisions: + +```text +SimplerEnv: 06accaca93535902d408da4855f21cece12bceb7 +ManiSkill2_real2sim: ef7a4d4fdf4b69f2c2154db5b15b9ac8dfe10682 +``` + +```bash +conda env create -f eval/simpler_env/environment.yaml +conda activate robotcpp-simpler-env + +git clone --recurse-submodules https://github.com/simpler-env/SimplerEnv \ + ckpts/simpler_env/source/SimplerEnv +git -C ckpts/simpler_env/source/SimplerEnv checkout \ + 06accaca93535902d408da4855f21cece12bceb7 +git -C ckpts/simpler_env/source/SimplerEnv submodule update --init --recursive + +pip install -e ckpts/simpler_env/source/SimplerEnv/ManiSkill2_real2sim +pip install -e ckpts/simpler_env/source/SimplerEnv +``` + +Headless simulation requires a working Vulkan ICD. Run SimplerEnv's environment +test first to confirm that SAPIEN can find a rendering device. + +## Run + +`VARIANT` accepts `oft`, `groot`, `pi_v3`, `qwen25_oft`, `qwen25_groot`, +`qwen25_pi`, and `qwen25_fast`. + +Run the full profile: + +```bash +CUDA_VISIBLE_DEVICES=0 \ +VARIANT=oft \ +OUTPUT=ckpts/starvla/results/oft/bridge.json \ +bash eval/simpler_env/scripts/run_model_server.sh +``` + +Run one smoke episode: + +```bash +CUDA_VISIBLE_DEVICES=0 \ +VARIANT=groot TASK_IDS=0 EPISODE_IDS=0 \ +bash eval/simpler_env/scripts/run_model_server.sh +``` + +The script reads three GGUF files from `ckpts/starvla/gguf/` and uses +`build_cuda/bin/model-server` by default. Common overrides are `GGUF_DIR`, +`SERVER_BIN`, `PYTHON`, `SIMPLER_ENV_ROOT`, `TASK_IDS`, +`EPISODE_IDS`, `REPEATS`, and `OUTPUT`. + +Each task/repeat starts a fresh model-server. Results include checkpoint +identity, rollout records, success rates, and timing summaries. + +## Latency + +Benchmark the official PyTorch checkpoint directly: + +```bash +CUDA_VISIBLE_DEVICES=0 python -m eval.simpler_env.runners.latency_starvla \ + --variant oft --compile-model +``` + +The runner selects the checkpoint, Qwen assets, and Bridge normalization from +the StarVLA catalog. It reports policy, action unnormalization, and total +latency after 5 warmup calls and 20 measured calls. Policy latency includes +StarVLA's image/text preprocessing and model forward. Omit `--compile-model` +for eager PyTorch. Compilation is lazy; the first FAST warmup can take several +minutes and is not included in the reported measurements. + +For the robot.cpp model-server path, use the common server benchmark: + +```bash +CUDA_VISIBLE_DEVICES=0 N_BATCH=2048 SKIP_BUILD=1 \ +GGUF_DIR="$PWD/ckpts/starvla/gguf/oft" \ +bash robot_server/test/test_server_latency.sh starvla linux-cuda starvla-bridge +``` diff --git a/eval/simpler_env/README_ZH.md b/eval/simpler_env/README_ZH.md new file mode 100644 index 0000000..4a049b2 --- /dev/null +++ b/eval/simpler_env/README_ZH.md @@ -0,0 +1,97 @@ +# SimplerEnv WidowX Bridge 评测 + +本目录使用 robot.cpp 的 StarVLA GGUF runtime 运行 SimplerEnv WidowX Bridge 任务。 + +## 评测设置 + +- 四个 Bridge 任务，每个任务包含 object episode `0..23` +- 每个 episode 最多 120 步，控制频率 5 Hz +- visual-matching RGB overlay 使用 OpenCV `INTER_AREA` 缩放到 224x224 +- 每步预测一个 action chunk，并对最近七次预测做自适应集成 + +完整评测包含 96 个 rollout。结果文件同时记录总体和各任务成功率；子集运行会标记为 +`partial` coverage。 + +## 安装 + +先按 [StarVLA 转换说明](../../tools/hf2gguf/starvla/README.md) 生成 GGUF，并完成 CUDA +构建。 + +SimplerEnv 使用以下 revision： + +```text +SimplerEnv: 06accaca93535902d408da4855f21cece12bceb7 +ManiSkill2_real2sim: ef7a4d4fdf4b69f2c2154db5b15b9ac8dfe10682 +``` + +```bash +conda env create -f eval/simpler_env/environment.yaml +conda activate robotcpp-simpler-env + +git clone --recurse-submodules https://github.com/simpler-env/SimplerEnv \ + ckpts/simpler_env/source/SimplerEnv +git -C ckpts/simpler_env/source/SimplerEnv checkout \ + 06accaca93535902d408da4855f21cece12bceb7 +git -C ckpts/simpler_env/source/SimplerEnv submodule update --init --recursive + +pip install -e ckpts/simpler_env/source/SimplerEnv/ManiSkill2_real2sim +pip install -e ckpts/simpler_env/source/SimplerEnv +``` + +无头运行需要可用的 Vulkan ICD。请先运行 SimplerEnv 自带的环境测试，确认 SAPIEN 能找到 +渲染设备。 + +## 运行 + +`VARIANT` 支持： + +```text +oft groot pi_v3 qwen25_oft qwen25_groot qwen25_pi qwen25_fast +``` + +完整运行： + +```bash +CUDA_VISIBLE_DEVICES=0 \ +VARIANT=oft \ +OUTPUT=ckpts/starvla/results/oft/bridge.json \ +bash eval/simpler_env/scripts/run_model_server.sh +``` + +快速检查一个 episode： + +```bash +CUDA_VISIBLE_DEVICES=0 \ +VARIANT=groot TASK_IDS=0 EPISODE_IDS=0 \ +bash eval/simpler_env/scripts/run_model_server.sh +``` + +脚本默认从 `ckpts/starvla/gguf/` 读取三个 GGUF，并使用 +`build_cuda/bin/model-server`。常用覆盖项包括 `GGUF_DIR`、`SERVER_BIN`、`PYTHON`、 +`SIMPLER_ENV_ROOT`、`TASK_IDS`、`EPISODE_IDS`、`REPEATS` 和 `OUTPUT`。 + +每个 task/repeat 会启动新的 model-server。结果包含 checkpoint 标识、rollout 明细、成功率 +和各阶段耗时。 + +## 延迟测试 + +直接测试官方 PyTorch checkpoint： + +```bash +CUDA_VISIBLE_DEVICES=0 python -m eval.simpler_env.runners.latency_starvla \ + --variant oft --compile-model +``` + +runner 会根据 StarVLA catalog 选择 checkpoint、Qwen 资源和 Bridge 归一化配置。默认先预热 +5 次，再统计 20 次推理，并分别报告 policy、action 反归一化和总耗时。policy 耗时包含 +StarVLA 的图像/文本预处理和模型 forward。去掉 `--compile-model` 即可测试 eager +PyTorch。`torch.compile` 为惰性编译；FAST 第一次预热可能需要几分钟，这部分不会计入 +最终统计。 + +robot.cpp model-server 使用统一的服务端测试脚本： + +```bash +CUDA_VISIBLE_DEVICES=0 N_BATCH=2048 SKIP_BUILD=1 \ +GGUF_DIR="$PWD/ckpts/starvla/gguf/oft" \ +bash robot_server/test/test_server_latency.sh starvla linux-cuda starvla-bridge +``` diff --git a/eval/simpler_env/__init__.py b/eval/simpler_env/__init__.py new file mode 100644 index 0000000..55ef27c --- /dev/null +++ b/eval/simpler_env/__init__.py @@ -0,0 +1 @@ +"""SimplerEnv evaluation integration.""" diff --git a/eval/simpler_env/environment.yaml b/eval/simpler_env/environment.yaml new file mode 100644 index 0000000..1b17b4d --- /dev/null +++ b/eval/simpler_env/environment.yaml @@ -0,0 +1,19 @@ +name: robotcpp-simpler-env +channels: + - conda-forge +dependencies: + - python=3.10 + - pip + - ffmpeg + - pip: + - numpy==1.24.4 + - scipy==1.11.4 + - opencv-python==4.11.0.86 + - opencv-python-headless==4.11.0.86 + - setuptools<81 + - transforms3d + - matplotlib + - mediapy + - tyro + - msgpack + - websockets diff --git a/eval/simpler_env/policy/__init__.py b/eval/simpler_env/policy/__init__.py new file mode 100644 index 0000000..3dbcd28 --- /dev/null +++ b/eval/simpler_env/policy/__init__.py @@ -0,0 +1 @@ +"""Policies used by the SimplerEnv runners.""" diff --git a/eval/simpler_env/policy/model_server.py b/eval/simpler_env/policy/model_server.py new file mode 100644 index 0000000..b5fb743 --- /dev/null +++ b/eval/simpler_env/policy/model_server.py @@ -0,0 +1,234 @@ +"""StarVLA model-server adapter for the SimplerEnv WidowX benchmark.""" + +from __future__ import annotations + +import time +from collections import deque +from typing import Any + +import numpy as np + +from eval.libero.policy.model_server import ServerTiming +from robot_client.python.model_client import ModelClient, ModelResponse + + +DEFAULT_IMAGE_NAME = "image_0" +DEFAULT_IMAGE_SIZE = (224, 224) +DEFAULT_ACTION_ENSEMBLE_HORIZON = 7 +DEFAULT_ADAPTIVE_ENSEMBLE_ALPHA = 0.1 + + +class AdaptiveEnsembler: + """StarVLA's cosine-similarity weighted temporal action ensemble.""" + + def __init__(self, horizon: int, alpha: float = DEFAULT_ADAPTIVE_ENSEMBLE_ALPHA): + if horizon <= 0: + raise ValueError("action ensemble horizon must be positive") + self.horizon = int(horizon) + self.alpha = float(alpha) + self._history: deque[np.ndarray] = deque(maxlen=self.horizon) + + def reset(self) -> None: + self._history.clear() + + def ensemble_action(self, action_chunk: np.ndarray) -> np.ndarray: + chunk = np.asarray(action_chunk) + if not np.issubdtype(chunk.dtype, np.floating): + chunk = chunk.astype(np.float32) + if chunk.ndim not in (1, 2): + raise ValueError(f"expected a 1D action or 2D action chunk, got shape={chunk.shape}") + if chunk.ndim == 2 and chunk.shape[0] < min(len(self._history) + 1, self.horizon): + raise ValueError("action chunk is shorter than the active ensemble history") + + self._history.append(chunk) + count = len(self._history) + if chunk.ndim == 1: + current_predictions = np.stack(tuple(self._history)) + else: + current_predictions = np.stack( + [prediction[index] for index, prediction in zip(range(count - 1, -1, -1), self._history)] + ) + + reference = current_predictions[-1] + dot = np.sum(current_predictions * reference, axis=1) + norms = np.linalg.norm(current_predictions, axis=1) * np.linalg.norm(reference) + cosine = dot / (norms + 1e-7) + weights = np.exp(self.alpha * cosine) + weights /= weights.sum() + return np.sum(weights[:, None] * current_predictions, axis=0) + + +def resize_image_area(image: np.ndarray, image_size: tuple[int, int]) -> np.ndarray: + """Match the official StarVLA SimplerEnv client's OpenCV INTER_AREA resize.""" + + array = np.asarray(image) + if array.ndim != 3 or array.shape[2] != 3: + raise ValueError(f"expected an HWC RGB image, got shape={array.shape}") + if array.dtype != np.uint8: + raise ValueError(f"expected a uint8 RGB image, got dtype={array.dtype}") + width, height = (int(image_size[0]), int(image_size[1])) + if width <= 0 or height <= 0: + raise ValueError("image dimensions must be positive") + if array.shape[:2] == (height, width): + return np.ascontiguousarray(array) + try: + import cv2 + except ImportError as exc: + raise RuntimeError("opencv-python-headless is required to resize SimplerEnv observations") from exc + return cv2.resize(array, (width, height), interpolation=cv2.INTER_AREA) + + +def euler_xyz_to_axis_angle(rotation_delta: np.ndarray) -> np.ndarray: + """Use the same static-XYZ Euler convention as StarVLA's official adapter.""" + + roll, pitch, yaw = np.asarray(rotation_delta, dtype=np.float64).reshape(3) * 0.5 + sr, cr = np.sin(roll), np.cos(roll) + sp, cp = np.sin(pitch), np.cos(pitch) + sy, cy = np.sin(yaw), np.cos(yaw) + quaternion = np.asarray( + [ + cr * cp * cy + sr * sp * sy, + sr * cp * cy - cr * sp * sy, + cr * sp * cy + sr * cp * sy, + cr * cp * sy - sr * sp * cy, + ], + dtype=np.float64, + ) + quaternion /= np.linalg.norm(quaternion) + vector_norm = float(np.linalg.norm(quaternion[1:])) + if vector_norm <= 1e-12: + return np.zeros(3, dtype=np.float64) + angle = 2.0 * np.arccos(np.clip(quaternion[0], -1.0, 1.0)) + return quaternion[1:] * (angle / vector_norm) + + +class SimplerEnvModelServerPolicy: + """Closed-loop WidowX policy matching StarVLA's official SimplerEnv adapter.""" + + def __init__( + self, + *, + host: str = "127.0.0.1", + port: int = 5555, + timeout: float | None = 120.0, + image_name: str = DEFAULT_IMAGE_NAME, + image_size: tuple[int, int] = DEFAULT_IMAGE_SIZE, + action_scale: float = 1.0, + action_ensemble: bool = True, + action_ensemble_horizon: int = DEFAULT_ACTION_ENSEMBLE_HORIZON, + adaptive_ensemble_alpha: float = DEFAULT_ADAPTIVE_ENSEMBLE_ALPHA, + client: ModelClient | None = None, + ): + self.client = client or ModelClient(host=host, port=port, timeout=timeout) + self.image_name = str(image_name) + self.image_size = (int(image_size[0]), int(image_size[1])) + self.action_scale = float(action_scale) + self.action_ensembler = ( + AdaptiveEnsembler(action_ensemble_horizon, adaptive_ensemble_alpha) + if action_ensemble + else None + ) + self.task_description: str | None = None + self.predict_calls = 0 + self.timing_records: list[ServerTiming] = [] + self._action_shape: tuple[int, int] | None = None + + def health(self) -> str: + return self.client.health() + + def action_shape(self) -> tuple[int, int]: + if self._action_shape is None: + raise RuntimeError("model-server has not returned an action chunk") + return self._action_shape + + def _validate_response_actions(self, response: ModelResponse) -> np.ndarray: + shape = (int(response.chunk_size), int(response.action_dim)) + if shape[0] <= 0 or shape[1] <= 0: + raise RuntimeError(f"model-server returned an invalid action shape: {shape}") + # Protocol actions are FP32. Preserve that dtype through temporal + # ensembling to match StarVLA's official SimplerEnv client. + actions = np.asarray(response.actions, dtype=np.float32) + if actions.shape != shape: + raise RuntimeError( + "model-server returned an invalid action matrix: " + f"wire_shape={shape}, decoded_shape={actions.shape}" + ) + if not np.isfinite(actions).all(): + raise RuntimeError("model-server returned non-finite action values") + if shape[1] != 7: + raise RuntimeError(f"SimplerEnv WidowX requires 7D actions, got {shape}") + if self.action_ensembler is not None and shape[0] < self.action_ensembler.horizon: + raise RuntimeError( + f"action chunk {shape[0]} is shorter than ensemble horizon " + f"{self.action_ensembler.horizon}" + ) + if self._action_shape is not None and shape != self._action_shape: + raise RuntimeError( + f"model-server action shape changed from {self._action_shape} to {shape}" + ) + self._action_shape = shape + return actions + + def reset( + self, + task_description: str | None = None, + *, + reset_server: bool = True, + ) -> None: + self.task_description = task_description + if self.action_ensembler is not None: + self.action_ensembler.reset() + if reset_server: + self.client.reset() + + def build_observation(self, image: np.ndarray, task_description: str) -> dict[str, Any]: + resized = resize_image_area(image, self.image_size) + observation = { + "images": [{"name": self.image_name, "image": resized}], + "state": [], + "prompt": task_description, + } + return observation + + def predict_action_chunk(self, image: np.ndarray, task_description: str) -> ModelResponse: + request = self.build_observation(image, task_description) + started = time.perf_counter() + response = self.client.predict(request) + self._validate_response_actions(response) + self.timing_records.append( + ServerTiming( + roundtrip_ms=(time.perf_counter() - started) * 1000.0, + timings=response.timings, + ) + ) + self.predict_calls += 1 + return response + + def step( + self, image: np.ndarray, task_description: str | None = None + ) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]: + if task_description is not None and task_description != self.task_description: + self.reset(task_description, reset_server=False) + if self.task_description is None: + raise ValueError("task_description must be set before policy.step") + + response = self.predict_action_chunk(image, self.task_description) + actions = np.asarray(response.actions, dtype=np.float32) + selected = ( + self.action_ensembler.ensemble_action(actions) + if self.action_ensembler is not None + else actions[0] + ) + + raw_action = { + "world_vector": selected[:3].copy(), + "rotation_delta": selected[3:6].copy(), + "open_gripper": selected[6:7].copy(), + } + action = { + "world_vector": raw_action["world_vector"] * self.action_scale, + "rot_axangle": euler_xyz_to_axis_angle(raw_action["rotation_delta"]) * self.action_scale, + "gripper": 2.0 * (raw_action["open_gripper"] > 0.5).astype(np.float64) - 1.0, + "terminate_episode": np.asarray([0.0], dtype=np.float64), + } + return raw_action, action diff --git a/eval/simpler_env/runners/__init__.py b/eval/simpler_env/runners/__init__.py new file mode 100644 index 0000000..8145c73 --- /dev/null +++ b/eval/simpler_env/runners/__init__.py @@ -0,0 +1 @@ +"""SimplerEnv evaluation runners.""" diff --git a/eval/simpler_env/runners/latency_starvla.py b/eval/simpler_env/runners/latency_starvla.py new file mode 100644 index 0000000..52972a7 --- /dev/null +++ b/eval/simpler_env/runners/latency_starvla.py @@ -0,0 +1,424 @@ +#!/usr/bin/env python3 + +from __future__ import annotations + +import argparse +import contextlib +import gc +import io +import json +import statistics +import subprocess +import sys +import tarfile +import tempfile +import time +from collections.abc import Iterator, Mapping +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from PIL import Image + +from eval.libero.utils.common import DEFAULT_RESULTS_DIR, timestamp, write_json + + +REPO_ROOT = Path(__file__).resolve().parents[3] +STARVLA_TOOLS = REPO_ROOT / "tools" / "hf2gguf" / "starvla" +if str(STARVLA_TOOLS) not in sys.path: + sys.path.insert(0, str(STARVLA_TOOLS)) + +from starvla_checkpoint import ( # noqa: E402 + DEFAULT_CATALOG, + get_qwen_asset, + get_variant, + load_catalog, + resolve_effective_config, +) + + +DEFAULT_PROMPT = "grab the block." +LEGACY_PI_REVISION = "e872a8579055f9332add8a2549b9fd5599e11510" +FULL_BF16_VARIANTS = {"oft", "qwen25_oft", "qwen25_pi"} +FRAMEWORK_CLASSES = { + "oft": ("starVLA.model.framework.VLM4A.QwenOFT", "Qwenvl_OFT"), + "groot": ("starVLA.model.framework.VLM4A.QwenGR00T", "Qwen_GR00T"), + "pi_v3": ("starVLA.model.framework.VLM4A.QwenPI_v3", "Qwen_PI_v3"), + "fast": ("starVLA.model.framework.VLM4A.QwenFast", "Qwenvl_Fast"), +} + + +def build_parser() -> argparse.ArgumentParser: + variants = tuple(load_catalog(DEFAULT_CATALOG)["variants"]) + parser = argparse.ArgumentParser(description="Benchmark an official StarVLA checkpoint with PyTorch.") + parser.add_argument("--variant", choices=variants, required=True) + parser.add_argument("--checkpoint-root", type=Path, default=REPO_ROOT / "ckpts" / "starvla") + parser.add_argument("--starvla-source", type=Path) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--compile-model", action=argparse.BooleanOptionalAction, default=False) + parser.add_argument("--compile-mode", default="default") + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--loops", type=int, default=20) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--prompt", default=DEFAULT_PROMPT) + parser.add_argument("--image-height", type=int, default=224) + parser.add_argument("--image-width", type=int, default=224) + parser.add_argument("--output", type=Path) + return parser + + +def percentile(values: list[float], pct: float) -> float: + if len(values) == 1: + return values[0] + ordered = sorted(values) + pos = (len(ordered) - 1) * pct / 100.0 + lo = int(pos) + hi = min(lo + 1, len(ordered) - 1) + return ordered[lo] * (hi - pos) + ordered[hi] * (pos - lo) + + +def summarize(values: list[float]) -> dict[str, float | int]: + return { + "count": len(values), + "avg": statistics.fmean(values), + "min": min(values), + "p50": percentile(values, 50), + "p90": percentile(values, 90), + "p99": percentile(values, 99), + "max": max(values), + } + + +def sync(device: str) -> None: + if device.startswith("cuda"): + torch.cuda.synchronize(torch.device(device)) + + +def checkpoint_paths(checkpoint_root: Path, variant_name: str) -> dict[str, Any]: + catalog = load_catalog(DEFAULT_CATALOG) + variant = get_variant(catalog, variant_name) + _qwen_name, qwen = get_qwen_asset(catalog, variant) + policy_dir = checkpoint_root / "sources" / variant["directory"] / variant["revision"] + qwen_dir = checkpoint_root / "sources" / qwen["directory"] / qwen["revision"] + checkpoint = policy_dir / variant["checkpoint"]["path"] + for path in (policy_dir, qwen_dir, checkpoint): + if not path.exists(): + raise FileNotFoundError(path) + return { + "catalog": catalog, + "variant": variant, + "policy_dir": policy_dir.resolve(), + "qwen_dir": qwen_dir.resolve(), + "checkpoint": checkpoint.resolve(), + } + + +def verify_source(source: Path, catalog: Mapping[str, Any]) -> str: + revision = subprocess.run( + ["git", "-C", str(source), "rev-parse", "HEAD"], + check=True, + capture_output=True, + text=True, + ).stdout.strip() + expected = catalog["source_revisions"]["starvla"] + if revision != expected: + raise RuntimeError(f"StarVLA source revision must be {expected}, got {revision}") + changes = subprocess.run( + ["git", "-C", str(source), "status", "--porcelain"], + check=True, + capture_output=True, + text=True, + ).stdout.strip() + if changes: + raise RuntimeError(f"StarVLA source checkout is not clean:\n{changes}") + return revision + + +@contextlib.contextmanager +def qwen_alias(qwen_dir: Path, backbone: str) -> Iterator[Path]: + name = "Qwen3-VL-4B-Instruct" if backbone == "qwen3_vl" else "Qwen2.5-VL-3B-Instruct" + with tempfile.TemporaryDirectory(prefix="starvla-latency-qwen-") as temporary: + alias = Path(temporary) / name + alias.symlink_to(qwen_dir, target_is_directory=True) + yield alias + + +@contextlib.contextmanager +def config_only_qwen(qwen_dir: Path, backbone: str) -> Iterator[None]: + import transformers + + model_class = ( + transformers.Qwen3VLForConditionalGeneration + if backbone == "qwen3_vl" + else transformers.Qwen2_5_VLForConditionalGeneration + ) + original = model_class.__dict__.get("from_pretrained") + + def from_config_only(_model_id: str, **_kwargs: Any) -> Any: + config = transformers.AutoConfig.from_pretrained(qwen_dir, local_files_only=True) + config._attn_implementation = "sdpa" + previous = torch.get_default_dtype() + try: + torch.set_default_dtype(torch.bfloat16) + with transformers.modeling_utils.no_init_weights(): + return model_class(config) + finally: + torch.set_default_dtype(previous) + + model_class.from_pretrained = staticmethod(from_config_only) + try: + yield + finally: + if original is None: + delattr(model_class, "from_pretrained") + else: + model_class.from_pretrained = original + + +def extract_legacy_source(source: Path) -> tempfile.TemporaryDirectory[str]: + archive = subprocess.run( + ["git", "-C", str(source), "archive", "--format=tar", LEGACY_PI_REVISION], + check=True, + stdout=subprocess.PIPE, + ).stdout + holder: tempfile.TemporaryDirectory[str] = tempfile.TemporaryDirectory(prefix="starvla-pi-") + with tarfile.open(fileobj=io.BytesIO(archive), mode="r:") as stream: + stream.extractall( + holder.name, + filter=lambda member, target: None + if member.issym() or member.islnk() + else tarfile.data_filter(member, target), + ) + return holder + + +def fast_config(policy_dir: Path) -> dict[str, Any]: + import yaml + + config = yaml.safe_load((policy_dir / "config.yaml").read_text(encoding="utf-8")) + config["framework"]["name"] = "QwenFast" + config["framework"]["action_model"]["action_model_type"] = "FAST" + config["framework"]["action_model"]["action_horizon"] = 16 + return config + + +def install_policy_dtype_bridge(framework: Any, framework_name: str) -> None: + if framework_name not in {"groot", "pi_v3"}: + return + action_model = framework.action_model + original_encoder = action_model.action_encoder.forward + original_dit = action_model.model.forward + + def encoder(actions: Any, timesteps: Any) -> Any: + return original_encoder(actions.float(), timesteps) + + def dit(*args: Any, **kwargs: Any) -> Any: + conditioning = kwargs.get("encoder_hidden_states", args[1] if len(args) > 1 else None) + if isinstance(conditioning, (list, tuple)): + conditioning = [value.float() for value in conditioning] + else: + conditioning = conditioning.float() + if "encoder_hidden_states" in kwargs: + kwargs["encoder_hidden_states"] = conditioning + else: + args = (args[0], conditioning, *args[2:]) + return original_dit(*args, **kwargs) + + action_model.action_encoder.forward = encoder + action_model.model.forward = dit + + +def load_framework(paths: Mapping[str, Any], source: Path, device: str) -> tuple[Any, Any]: + import importlib + + import yaml + + variant = paths["variant"] + variant_name = variant["_catalog_key"] + framework_name = variant["framework"] + runtime_source = source + holder = None + if variant_name == "qwen25_pi": + holder = extract_legacy_source(source) + runtime_source = Path(holder.name) + + sys.path.insert(0, str(runtime_source)) + try: + if variant_name == "qwen25_pi": + module_name, class_name = "starVLA.model.framework.QwenPI", "Qwen_PI" + config = yaml.safe_load((paths["policy_dir"] / "config.yaml").read_text(encoding="utf-8")) + else: + module_name, class_name = FRAMEWORK_CLASSES[framework_name] + config = ( + fast_config(paths["policy_dir"]) + if framework_name == "fast" + else resolve_effective_config(paths["policy_dir"], variant_name, variant) + ) + + from starVLA.model.framework import share_tools + + with qwen_alias(paths["qwen_dir"], variant["backbone"]) as alias: + config["framework"]["qwenvl"]["base_vlm"] = str(alias) + config["framework"]["qwenvl"]["attn_implementation"] = "sdpa" + cfg = share_tools.dict_to_namespace(config) + cfg.trainer.pretrained_checkpoint = None + module = importlib.import_module(module_name) + if framework_name == "fast": + from starVLA.model.modules.action_model.fast_ActionHeader import Fast_Action_Tokenizer + + codec = paths["catalog"]["shared_assets"]["fast_codec"] + codec_dir = ( + paths["policy_dir"].parents[1] / codec["directory"] / codec["revision"] + ) + module.get_action_model = lambda config=None: Fast_Action_Tokenizer(str(codec_dir)) + with config_only_qwen(paths["qwen_dir"], variant["backbone"]): + framework = getattr(module, class_name)(cfg) + + state = torch.load(paths["checkpoint"], map_location="cpu", mmap=True, weights_only=True) + framework.load_state_dict(state, strict=True) + del state + gc.collect() + framework.norm_stats = json.loads( + (paths["policy_dir"] / "dataset_statistics.json").read_text(encoding="utf-8") + ) + if variant_name in FULL_BF16_VARIANTS: + framework = framework.to(dtype=torch.bfloat16) + framework = framework.to(device).eval() + install_policy_dtype_bridge(framework, framework_name) + return framework, holder + except Exception: + if holder is not None: + holder.cleanup() + raise + finally: + if sys.path and sys.path[0] == str(runtime_source): + del sys.path[0] + + +def enable_compile(framework: Any, backbone: str, framework_name: str, mode: str) -> None: + qwen = framework.qwen_vl_interface + if framework_name == "fast": + qwen.model.forward = torch.compile(qwen.model.forward, mode=mode, fullgraph=False) + elif backbone == "qwen3_vl": + model = qwen.model.model + model.visual.forward = torch.compile(model.visual.forward, mode=mode, fullgraph=False) + for layer in model.language_model.layers: + layer.forward = torch.compile(layer.forward, mode=mode, fullgraph=False) + else: + qwen.forward = torch.compile(qwen.forward, mode=mode, fullgraph=False) + + if framework_name == "oft": + framework.action_model.predict_action = torch.compile( + framework.action_model.predict_action, mode=mode, fullgraph=False + ) + elif framework_name != "fast": + framework.action_model.model.forward = torch.compile( + framework.action_model.model.forward, mode=mode, fullgraph=False + ) + + +def unnormalize(normalized: Any, statistics: Mapping[str, Any]) -> np.ndarray: + profile_name = next(iter(statistics)) + stats = statistics[profile_name]["action"] + values = np.asarray(normalized, dtype=np.float32) + if values.shape != (1, 16, 7) or not np.isfinite(values).all(): + raise ValueError(f"StarVLA returned invalid normalized actions: {values.shape}") + q01 = np.asarray(stats["q01"], dtype=np.float32) + q99 = np.asarray(stats["q99"], dtype=np.float32) + mask = np.asarray(stats["mask"], dtype=np.bool_) + result = np.empty_like(values) + result[..., mask] = (values[..., mask] + 1.0) * 0.5 * (q99[mask] - q01[mask]) + q01[mask] + result[..., ~mask] = (values[..., ~mask] > 0.5).astype(np.float32) + return result + + +def predict(framework: Any, variant: str, image: Image.Image, prompt: str) -> Mapping[str, Any]: + if variant == "qwen25_pi": + return framework.predict_action(batch_images=[[image]], instructions=[prompt], state=None) + return framework.predict_action(examples=[{"image": [image], "lang": prompt}]) + + +def main() -> int: + args = build_parser().parse_args() + if args.warmup < 0 or args.loops <= 0: + raise ValueError("--warmup must be non-negative and --loops must be positive") + if not args.device.startswith("cuda") or not torch.cuda.is_available(): + raise RuntimeError("StarVLA latency currently requires CUDA") + + torch.manual_seed(args.seed) + np.random.seed(args.seed) + checkpoint_root = args.checkpoint_root.resolve() + source = (args.starvla_source or checkpoint_root / "source" / "starvla").resolve() + paths = checkpoint_paths(checkpoint_root, args.variant) + source_revision = verify_source(source, paths["catalog"]) + output = args.output or DEFAULT_RESULTS_DIR / f"starvla-policy-latency-{args.variant}-{timestamp()}.json" + + load_start = time.perf_counter() + framework, holder = load_framework(paths, source, args.device) + if args.compile_model: + enable_compile(framework, paths["variant"]["backbone"], paths["variant"]["framework"], args.compile_mode) + load_ms = (time.perf_counter() - load_start) * 1000.0 + + rng = np.random.default_rng(args.seed) + image = Image.fromarray( + rng.integers(0, 256, size=(args.image_height, args.image_width, 3), dtype=np.uint8), mode="RGB" + ) + rows: list[dict[str, float]] = [] + actions = None + print( + f"StarVLA latency: variant={args.variant} warmup={args.warmup} loops={args.loops} " + f"compile_model={args.compile_model} device={args.device}" + ) + for index in range(args.warmup + args.loops): + sync(args.device) + started = time.perf_counter() + output_value = predict(framework, args.variant, image, args.prompt) + sync(args.device) + policy_ms = (time.perf_counter() - started) * 1000.0 + + unnorm_started = time.perf_counter() + actions = unnormalize(output_value["normalized_actions"], framework.norm_stats) + unnormalize_ms = (time.perf_counter() - unnorm_started) * 1000.0 + total_ms = policy_ms + unnormalize_ms + if index >= args.warmup: + rows.append({"policy_ms": policy_ms, "unnormalize_ms": unnormalize_ms, "total_ms": total_ms}) + print( + f"iter={index} policy_ms={policy_ms:.3f} unnormalize_ms={unnormalize_ms:.3f} " + f"total_ms={total_ms:.3f}", + flush=True, + ) + + assert actions is not None + variant = paths["variant"] + payload = { + "runner": "starvla-policy-latency", + "variant": args.variant, + "framework": variant["framework"], + "backbone": variant["backbone"], + "checkpoint": {"repo_id": variant["repo_id"], "revision": variant["revision"]}, + "starvla_revision": source_revision, + "device": args.device, + "compile_model": args.compile_model, + "compile_mode": args.compile_mode if args.compile_model else None, + "warmup": args.warmup, + "loops": args.loops, + "load_ms": load_ms, + "action_shape": list(actions.shape), + "raw_input": {"image_shape_hwc": [args.image_height, args.image_width, 3], "prompt": args.prompt}, + "timing_ms": { + key: summarize([row[key] for row in rows]) + for key in ("policy_ms", "unnormalize_ms", "total_ms") + }, + "rows": rows, + } + write_json(output, payload) + print(f"wrote {output}") + print(json.dumps(payload["timing_ms"], indent=2)) + if holder is not None: + holder.cleanup() + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/eval/simpler_env/runners/run_model_server.py b/eval/simpler_env/runners/run_model_server.py new file mode 100755 index 0000000..74dcd8c --- /dev/null +++ b/eval/simpler_env/runners/run_model_server.py @@ -0,0 +1,404 @@ +#!/usr/bin/env python3 + +from __future__ import annotations + +import argparse +import time +from collections import defaultdict +from copy import copy +from pathlib import Path +from typing import Any + +import numpy as np + +from eval.libero.policy.model_server import ( + average_timing, + maybe_launch_server, + parse_server_env, + server_command, + stop_server, + timing_summary, +) +from eval.libero.utils.common import DEFAULT_RESULTS_DIR, aggregate_episodes, timestamp, write_json +from eval.simpler_env.policy.model_server import ( + DEFAULT_ACTION_ENSEMBLE_HORIZON, + DEFAULT_ADAPTIVE_ENSEMBLE_ALPHA, + DEFAULT_IMAGE_NAME, + DEFAULT_IMAGE_SIZE, + SimplerEnvModelServerPolicy, +) +from eval.simpler_env.utils.environment import ( + BRIDGE_EPISODE_COUNT, + BRIDGE_SUITE, + BRIDGE_TASKS, + BridgeTask, + apply_runtime_env, + close_env, + language_instruction, + make_env, + observation_image, + parse_episode_ids, + parse_task_ids, + reset_env, + selected_tasks, + simpler_env_root, +) + + +def _first_bool(value: Any) -> bool: + array = np.asarray(value).reshape(-1) + if array.size != 1: + raise RuntimeError(f"expected one termination value, got shape={np.asarray(value).shape}") + return bool(array[0]) + + +def _first_float(value: Any) -> float: + array = np.asarray(value).reshape(-1) + if array.size != 1 or not np.isfinite(array[0]): + raise RuntimeError(f"expected one finite reward value, got {value!r}") + return float(array[0]) + + +def _write_video(path: Path, images: list[np.ndarray], fps: int) -> None: + try: + from simpler_env.utils.visualization import write_video + except ImportError as exc: + raise RuntimeError("failed to import SimplerEnv video writer") from exc + path.parent.mkdir(parents=True, exist_ok=True) + write_video(str(path), images, fps=fps) + + +def _command_with_noise_seed(command: list[str], seed: int) -> list[str]: + result = list(command) + for index, value in enumerate(result): + if value == "--noise-seed": + result[index + 1] = str(seed) + return result + if value.startswith("--noise-seed="): + result[index] = f"--noise-seed={seed}" + return result + return [*result, "--noise-seed", str(seed)] + + +def _launch_fresh_server(args: argparse.Namespace, policy: SimplerEnvModelServerPolicy): + try: + health = policy.health() + except OSError: + pass + else: + raise RuntimeError( + f"refusing to reuse model-server at {args.host}:{args.port}: {health}" + ) + process = maybe_launch_server(args, policy) + if process is None: + raise RuntimeError("model-server launch did not create a process") + return process + + +def model_record( + args: argparse.Namespace, policy: SimplerEnvModelServerPolicy +) -> dict[str, Any]: + chunk_size, action_dim = policy.action_shape() + return { + "model_type": args.expected_model_type, + "variant": args.variant, + "framework": args.expected_framework, + "checkpoint_revision": args.expected_checkpoint_revision, + "checkpoint_sha256": args.expected_checkpoint_sha256, + "qwen_revision": args.expected_qwen_revision, + "starvla_revision": args.expected_starvla_revision, + "chunk_size": chunk_size, + "action_dim": action_dim, + } + + +def aggregate_task_repeats(episodes: list[dict[str, Any]]) -> list[dict[str, Any]]: + groups: dict[tuple[int, int], list[dict[str, Any]]] = defaultdict(list) + for episode in episodes: + groups[(int(episode["task_id"]), int(episode["repeat"]))].append(episode) + return [ + { + "suite": BRIDGE_SUITE, + "task_id": task_id, + "repeat": repeat, + **aggregate_episodes(rows)["overall"], + } + for (task_id, repeat), rows in sorted(groups.items()) + ] + + +def run_episode( + env: Any, + policy: SimplerEnvModelServerPolicy, + task_spec: BridgeTask, + episode_id: int, + *, + repeat: int, + max_episode_steps: int, + camera_name: str | None, + video_path: Path | None, + video_fps: int, +) -> dict[str, Any]: + observation, _ = reset_env(env, task_spec, episode_id) + task = language_instruction(env) + if task != task_spec.instruction: + raise RuntimeError( + f"unexpected instruction for {task_spec.env_name}: {task!r}" + ) + policy.reset(task, reset_server=True) + image = observation_image(env, observation, camera_name) + frames = [image] if video_path is not None else [] + start_predict_calls = policy.predict_calls + start_timing_index = len(policy.timing_records) + started = time.perf_counter() + rewards: list[float] = [] + success = False + terminated = False + truncated = False + steps = 0 + + while steps < max_episode_steps and not truncated: + _, action = policy.step(image, task) + env_action = np.concatenate( + [action["world_vector"], action["rot_axangle"], action["gripper"]] + ) + observation, reward, terminated_value, truncated_value, _ = env.step(env_action) + terminated = _first_bool(terminated_value) + truncated = _first_bool(truncated_value) + success = terminated + rewards.append(_first_float(reward)) + steps += 1 + if terminated or truncated: + break + task = language_instruction(env) + image = observation_image(env, observation, camera_name) + if frames: + frames.append(image) + + if video_path is not None: + _write_video(video_path, frames, video_fps) + records = policy.timing_records[start_timing_index:] + return { + "episode": int(episode_id), + "repeat": int(repeat), + "task": task_spec.instruction, + "task_name": task_spec.name, + "env_name": task_spec.env_name, + "success": bool(success), + "terminated": bool(terminated), + "truncated": bool(truncated), + "sum_reward": float(sum(rewards)), + "max_reward": float(max(rewards) if rewards else 0.0), + "steps": steps, + "elapsed_s": time.perf_counter() - started, + "predict_calls": policy.predict_calls - start_predict_calls, + "server_timing_avg_ms": average_timing(records), + "video": str(video_path) if video_path is not None else None, + } + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Evaluate model-server on SimplerEnv Bridge") + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, default=5555) + parser.add_argument("--launch-server", action="store_true") + parser.add_argument("--server-command", nargs=argparse.REMAINDER, help="must be last") + parser.add_argument("--server-env", action="append") + parser.add_argument("--server-wait-s", type=float, default=180.0) + parser.add_argument("--server-noise-seed-base", type=int, default=0) + parser.add_argument("--variant") + parser.add_argument("--expected-model-type") + parser.add_argument("--expected-checkpoint-revision") + parser.add_argument("--expected-checkpoint-sha256") + parser.add_argument("--expected-qwen-revision") + parser.add_argument("--expected-starvla-revision") + parser.add_argument("--expected-framework") + parser.add_argument("--task-ids", default="all") + parser.add_argument("--episode-ids", default="0:24") + parser.add_argument("--repeats", type=int, default=1) + parser.add_argument("--max-episode-steps", type=int, default=120) + parser.add_argument("--control-freq", type=int, default=5) + parser.add_argument("--sim-freq", type=int, default=500) + parser.add_argument("--image-name", default=DEFAULT_IMAGE_NAME) + parser.add_argument("--image-size", type=int, nargs=2, default=list(DEFAULT_IMAGE_SIZE)) + parser.add_argument("--action-scale", type=float, default=1.0) + parser.add_argument("--no-action-ensemble", action="store_true") + parser.add_argument( + "--action-ensemble-horizon", type=int, default=DEFAULT_ACTION_ENSEMBLE_HORIZON + ) + parser.add_argument( + "--adaptive-ensemble-alpha", type=float, default=DEFAULT_ADAPTIVE_ENSEMBLE_ALPHA + ) + parser.add_argument("--camera-name") + parser.add_argument("--no-rgb-overlay", action="store_true") + parser.add_argument("--enable-raytracing", action="store_true") + parser.add_argument("--simpler-env-root", type=Path) + parser.add_argument("--record-video", action="store_true") + parser.add_argument("--video-dir", type=Path) + parser.add_argument("--video-fps", type=int, default=5) + parser.add_argument("--output", type=Path) + return parser.parse_args(argv) + + +def _validate_args(args: argparse.Namespace) -> tuple[list[int], list[int]]: + positive = ( + args.repeats, + args.max_episode_steps, + args.control_freq, + args.sim_freq, + args.video_fps, + args.action_ensemble_horizon, + ) + if any(value <= 0 for value in positive): + raise ValueError("repeat, episode, timing, and ensemble values must be positive") + if args.server_noise_seed_base < 0: + raise ValueError("--server-noise-seed-base must be non-negative") + return parse_task_ids(args.task_ids), parse_episode_ids(args.episode_ids) + + +def run(args: argparse.Namespace) -> dict[str, Any]: + task_ids, episode_ids = _validate_args(args) + output = args.output or DEFAULT_RESULTS_DIR / f"server-simpler-env-bridge-{timestamp()}.json" + video_dir = args.video_dir or output.with_suffix("").with_name(output.stem + "-videos") + apply_runtime_env() + root = simpler_env_root(args.simpler_env_root) + policy = SimplerEnvModelServerPolicy( + host=args.host, + port=args.port, + image_name=args.image_name, + image_size=tuple(args.image_size), + action_scale=args.action_scale, + action_ensemble=not args.no_action_ensemble, + action_ensemble_horizon=args.action_ensemble_horizon, + adaptive_ensemble_alpha=args.adaptive_ensemble_alpha, + ) + episodes: list[dict[str, Any]] = [] + launches: list[dict[str, Any]] = [] + recorded_model: dict[str, Any] | None = None + process = None + + try: + if not args.launch_server: + maybe_launch_server(args, policy) + base_command = server_command(args) if args.launch_server else [] + for task_spec in selected_tasks(task_ids): + for repeat in range(1, args.repeats + 1): + derived_seed = args.server_noise_seed_base + task_spec.task_id * args.repeats + repeat - 1 + noise_seed = None + if args.launch_server: + noise_seed = derived_seed + launch_args = copy(args) + launch_args.server_command = _command_with_noise_seed(base_command, noise_seed) + process = _launch_fresh_server(launch_args, policy) + launches.append( + { + "task_id": task_spec.task_id, + "repeat": repeat, + "noise_seed": noise_seed, + } + ) + for episode_id in episode_ids: + video_path = ( + video_dir + / f"task-{task_spec.task_id}-{task_spec.name}" + / f"repeat-{repeat:02d}-episode-{episode_id:02d}.mp4" + if args.record_video + else None + ) + env = make_env( + task_spec, + root=root, + control_freq=args.control_freq, + sim_freq=args.sim_freq, + max_episode_steps=args.max_episode_steps, + use_rgb_overlay=not args.no_rgb_overlay, + enable_raytracing=args.enable_raytracing, + ) + try: + result = run_episode( + env, + policy, + task_spec, + episode_id, + repeat=repeat, + max_episode_steps=args.max_episode_steps, + camera_name=args.camera_name, + video_path=video_path, + video_fps=args.video_fps, + ) + finally: + close_env(env) + result.update( + suite=BRIDGE_SUITE, + task_id=task_spec.task_id, + noise_seed=noise_seed, + ) + episodes.append(result) + print( + f"bridge[{task_spec.task_id}] repeat={repeat} episode={episode_id} " + f"success={result['success']} steps={result['steps']}" + ) + current_model = model_record(args, policy) + if recorded_model is None: + recorded_model = current_model + elif current_model != recorded_model: + raise RuntimeError("model action contract changed between repeats") + if process is not None: + stop_server(process, policy) + process = None + finally: + stop_server(process, policy) + + assert recorded_model is not None + full_coverage = ( + task_ids == [task.task_id for task in BRIDGE_TASKS] + and episode_ids == list(range(BRIDGE_EPISODE_COUNT)) + ) + payload = { + "runner": "model-server", + "benchmark": { + "name": "SimplerEnv WidowX Bridge", + "suite": BRIDGE_SUITE, + "coverage": "full" if full_coverage else "partial", + }, + "config": { + "task_ids": task_ids, + "episode_ids": episode_ids, + "repeats": args.repeats, + "max_episode_steps": args.max_episode_steps, + "control_freq": args.control_freq, + "sim_freq": args.sim_freq, + "host": args.host, + "port": args.port, + "server_command": base_command or None, + "server_env": parse_server_env(args.server_env), + "server_launches": launches, + "image_name": args.image_name, + "image_size": args.image_size, + "action_scale": args.action_scale, + "action_ensemble": not args.no_action_ensemble, + "action_ensemble_horizon": args.action_ensemble_horizon, + "adaptive_ensemble_alpha": args.adaptive_ensemble_alpha, + "rgb_overlay": not args.no_rgb_overlay, + "camera_name": args.camera_name, + "raytracing": args.enable_raytracing, + }, + "model": recorded_model, + "episodes": episodes, + "per_task_repeat": aggregate_task_repeats(episodes), + "timing_ms": timing_summary(policy.timing_records), + **aggregate_episodes(episodes), + } + write_json(output, payload) + print(f"wrote {output}") + print(f"overall: {payload['overall']}") + return payload + + +def main() -> int: + run(parse_args()) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/eval/simpler_env/scripts/run_model_server.sh b/eval/simpler_env/scripts/run_model_server.sh new file mode 100755 index 0000000..1d3768a --- /dev/null +++ b/eval/simpler_env/scripts/run_model_server.sh @@ -0,0 +1,81 @@ +#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)" +REPO_ROOT="$(cd -- "${SCRIPT_DIR}/../../.." >/dev/null 2>&1 && pwd)" +cd "${REPO_ROOT}" + +source "${REPO_ROOT}/tools/hf2gguf/starvla/starvla_variant_config.sh" +VARIANT="${VARIANT:-groot}" +load_starvla_variant "${VARIANT}" +GGUF_DIR="${GGUF_DIR:-ckpts/starvla/gguf/${VARIANT}}" +LLM_GGUF="${LLM_GGUF:-${GGUF_DIR}/qwen-${ARTIFACT_STEM}-bf16.gguf}" +MMPROJ_GGUF="${MMPROJ_GGUF:-${GGUF_DIR}/mmproj-${ARTIFACT_STEM}-bf16.gguf}" +if [[ "${VARIANT}" == "qwen25_fast" ]]; then + POLICY_GGUF="${POLICY_GGUF:-${GGUF_DIR}/policy-qwen25-fast.gguf}" +else + POLICY_GGUF="${POLICY_GGUF:-${GGUF_DIR}/starvla-${ARTIFACT_STEM}-policy-fp32.gguf}" +fi + +BACKEND="${BACKEND:-linux-cuda}" +case "${BACKEND}" in + linux-cuda) BUILD_DIR="${BUILD_DIR:-${REPO_ROOT}/build_cuda}" ;; + linux-cpu) BUILD_DIR="${BUILD_DIR:-${REPO_ROOT}/build}" ;; + *) echo "unsupported BACKEND=${BACKEND}; expected linux-cuda or linux-cpu" >&2; exit 2 ;; +esac +SERVER_BIN="${SERVER_BIN:-${BUILD_DIR}/bin/model-server}" +PYTHON_BIN="${PYTHON:-ckpts/simpler_env/.venv/bin/python}" +HOST="${HOST:-127.0.0.1}" +PORT="${PORT:-5555}" +NOISE_SEED_BASE="${NOISE_SEED_BASE:-1000}" + +for path in "${LLM_GGUF}" "${MMPROJ_GGUF}" "${POLICY_GGUF}"; do + if [[ ! -f "${path}" ]]; then + echo "missing GGUF: ${path}" >&2 + exit 2 + fi +done +if [[ ! -x "${SERVER_BIN}" ]]; then + echo "model-server was not found or is not executable: ${SERVER_BIN}" >&2 + exit 2 +fi +if [[ ! -x "${PYTHON_BIN}" ]]; then + echo "SimplerEnv Python was not found: ${PYTHON_BIN}" >&2 + echo "follow eval/simpler_env/README_ZH.md to create it" >&2 + exit 2 +fi + +eval_cmd=( + "${PYTHON_BIN}" -m eval.simpler_env.runners.run_model_server + --launch-server + --host "${HOST}" + --port "${PORT}" + --variant "${VARIANT}" + --expected-model-type "${MODEL_TYPE}" + --expected-checkpoint-revision "${CHECKPOINT_REVISION}" + --expected-checkpoint-sha256 "${CHECKPOINT_SHA256}" + --expected-qwen-revision "${QWEN_REVISION}" + --expected-starvla-revision "${STARVLA_REVISION}" + --expected-framework "${FRAMEWORK}" + --server-noise-seed-base "${NOISE_SEED_BASE}" +) +[[ -n "${TASK_IDS:-}" ]] && eval_cmd+=(--task-ids "${TASK_IDS}") +[[ -n "${EPISODE_IDS:-}" ]] && eval_cmd+=(--episode-ids "${EPISODE_IDS}") +[[ -n "${REPEATS:-}" ]] && eval_cmd+=(--repeats "${REPEATS}") +[[ -n "${OUTPUT:-}" ]] && eval_cmd+=(--output "${OUTPUT}") +[[ -n "${SIMPLER_ENV_ROOT:-}" ]] && eval_cmd+=(--simpler-env-root "${SIMPLER_ENV_ROOT}") +[[ "${RECORD_VIDEO:-0}" == "1" ]] && eval_cmd+=(--record-video) +eval_cmd+=("$@") +eval_cmd+=( + --server-command + "${SERVER_BIN}" + --model-type "${MODEL_TYPE}" + --policy "${POLICY_GGUF}" + --llm "${LLM_GGUF}" + --mmproj "${MMPROJ_GGUF}" + --host "${HOST}" + --port "${PORT}" + --noise-seed "${NOISE_SEED_BASE}" +) + +exec "${eval_cmd[@]}" diff --git a/eval/simpler_env/utils/__init__.py b/eval/simpler_env/utils/__init__.py new file mode 100644 index 0000000..ea81195 --- /dev/null +++ b/eval/simpler_env/utils/__init__.py @@ -0,0 +1 @@ +"""SimplerEnv evaluation helpers.""" diff --git a/eval/simpler_env/utils/environment.py b/eval/simpler_env/utils/environment.py new file mode 100644 index 0000000..4978636 --- /dev/null +++ b/eval/simpler_env/utils/environment.py @@ -0,0 +1,230 @@ +"""Official StarVLA WidowX Bridge task and environment configuration.""" + +from __future__ import annotations + +import json +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np + + +BRIDGE_SUITE = "simpler_env_widowx_bridge" +BRIDGE_EPISODE_COUNT = 24 +BRIDGE_OFFICIAL_REPEATS = 4 +BRIDGE_CONTROL_MODE = "arm_pd_ee_target_delta_pose_align2_gripper_pd_joint_pos" + + +@dataclass(frozen=True) +class BridgeTask: + task_id: int + name: str + env_name: str + instruction: str + scene_name: str + robot: str + overlay_filename: str + robot_init_x: float + robot_init_y: float + + +BRIDGE_TASKS = ( + BridgeTask( + 0, + "stack_green_cube_on_yellow_cube", + "StackGreenCubeOnYellowCubeBakedTexInScene-v0", + "stack the green block on the yellow block", + "bridge_table_1_v1", + "widowx", + "bridge_real_eval_1.png", + 0.147, + 0.028, + ), + BridgeTask( + 1, + "put_carrot_on_plate", + "PutCarrotOnPlateInScene-v0", + "put carrot on plate", + "bridge_table_1_v1", + "widowx", + "bridge_real_eval_1.png", + 0.147, + 0.028, + ), + BridgeTask( + 2, + "put_spoon_on_table_cloth", + "PutSpoonOnTableClothInScene-v0", + "put the spoon on the towel", + "bridge_table_1_v1", + "widowx", + "bridge_real_eval_1.png", + 0.147, + 0.028, + ), + BridgeTask( + 3, + "put_eggplant_in_basket", + "PutEggplantInBasketScene-v0", + "put eggplant into yellow basket", + "bridge_table_1_v2", + "widowx_sink_camera_setup", + "bridge_sink.png", + 0.127, + 0.060, + ), +) + + +def parse_task_ids(value: str | None) -> list[int]: + if value is None or value.strip().lower() in {"", "all"}: + return [task.task_id for task in BRIDGE_TASKS] + text = value.strip() + decoded = json.loads(text) if text.startswith("[") else text.split(",") + if not isinstance(decoded, list): + raise ValueError("--task-ids must be 'all', a comma list, or a JSON list") + task_ids = [int(item) for item in decoded] + known = {task.task_id for task in BRIDGE_TASKS} + if len(set(task_ids)) != len(task_ids) or any(task_id not in known for task_id in task_ids): + raise ValueError(f"--task-ids must contain unique values from {sorted(known)}") + return task_ids + + +def selected_tasks(task_ids: list[int]) -> list[BridgeTask]: + by_id = {task.task_id: task for task in BRIDGE_TASKS} + return [by_id[task_id] for task_id in task_ids] + + +def parse_episode_ids(value: str | None) -> list[int]: + text = (value or "0:24").strip() + if ":" in text and not text.startswith("["): + fields = text.split(":") + if len(fields) not in (2, 3): + raise ValueError("--episode-ids range must be START:STOP or START:STOP:STEP") + start, stop = int(fields[0]), int(fields[1]) + step = int(fields[2]) if len(fields) == 3 else 1 + if step <= 0: + raise ValueError("--episode-ids range step must be positive") + episode_ids = list(range(start, stop, step)) + else: + decoded = json.loads(text) if text.startswith("[") else text.split(",") + if not isinstance(decoded, list): + raise ValueError("--episode-ids must be a comma list, JSON list, or range") + episode_ids = [int(item) for item in decoded if str(item).strip()] + if not episode_ids: + raise ValueError("--episode-ids must select at least one episode") + if len(set(episode_ids)) != len(episode_ids): + raise ValueError("--episode-ids must not contain duplicates") + if any(episode_id < 0 or episode_id >= BRIDGE_EPISODE_COUNT for episode_id in episode_ids): + raise ValueError(f"Bridge episode ids must be in [0, {BRIDGE_EPISODE_COUNT})") + return episode_ids + + +def apply_runtime_env() -> None: + os.environ["DISPLAY"] = "" + os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false") + + +def simpler_env_root(explicit_root: Path | None = None) -> Path: + try: + import mani_skill2_real2sim + import simpler_env + except ImportError as exc: + raise RuntimeError( + "simpler_env is not installed; follow eval/simpler_env/README_ZH.md" + ) from exc + + if explicit_root is not None: + root = explicit_root.expanduser().resolve() + else: + root = Path(simpler_env.__file__).resolve().parent.parent + if not (root / "ManiSkill2_real2sim").exists(): + raise RuntimeError(f"invalid SimplerEnv root (ManiSkill2_real2sim missing): {root}") + installed_simpler_root = Path(simpler_env.__file__).resolve().parent.parent + installed_maniskill_root = Path(mani_skill2_real2sim.__file__).resolve().parent.parent + expected_maniskill_root = (root / "ManiSkill2_real2sim").resolve() + if installed_simpler_root != root: + raise RuntimeError( + f"installed simpler_env comes from {installed_simpler_root}, expected {root}" + ) + if installed_maniskill_root != expected_maniskill_root: + raise RuntimeError( + "installed mani_skill2_real2sim comes from " + f"{installed_maniskill_root}, expected {expected_maniskill_root}" + ) + return root + + +def overlay_path(root: Path, task: BridgeTask) -> Path: + path = root / "ManiSkill2_real2sim" / "data" / "real_inpainting" / task.overlay_filename + if not path.is_file(): + raise RuntimeError(f"official Bridge RGB overlay is missing: {path}") + return path + + +def make_env( + task: BridgeTask, + *, + root: Path, + control_freq: int = 5, + sim_freq: int = 500, + max_episode_steps: int = 120, + use_rgb_overlay: bool = True, + enable_raytracing: bool = False, +) -> Any: + try: + from simpler_env.utils.env.env_builder import build_maniskill2_env + except ImportError as exc: + raise RuntimeError("failed to import the installed SimplerEnv environment builder") from exc + + additional: dict[str, Any] = {"shader_dir": "rt"} if enable_raytracing else {} + return build_maniskill2_env( + task.env_name, + **additional, + obs_mode="rgbd", + robot=task.robot, + sim_freq=int(sim_freq), + control_mode=BRIDGE_CONTROL_MODE, + control_freq=int(control_freq), + max_episode_steps=int(max_episode_steps), + scene_name=task.scene_name, + camera_cfgs={"add_segmentation": True}, + rgb_overlay_path=str(overlay_path(root, task)) if use_rgb_overlay else None, + ) + + +def reset_env(env: Any, task: BridgeTask, episode_id: int) -> tuple[Any, Any]: + options = { + "robot_init_options": { + "init_xy": np.asarray([task.robot_init_x, task.robot_init_y], dtype=np.float64), + "init_rot_quat": np.asarray([0.0, 0.0, 0.0, 1.0], dtype=np.float64), + }, + "obj_init_options": {"episode_id": int(episode_id)}, + } + return env.reset(options=options) + + +def observation_image(env: Any, observation: dict[str, Any], camera_name: str | None = None) -> np.ndarray: + try: + from simpler_env.utils.env.observation_utils import get_image_from_maniskill2_obs_dict + except ImportError as exc: + raise RuntimeError("failed to import SimplerEnv observation helpers") from exc + image = np.asarray(get_image_from_maniskill2_obs_dict(env, observation, camera_name=camera_name)) + if image.dtype != np.uint8 or image.ndim != 3 or image.shape[2] != 3: + raise RuntimeError(f"SimplerEnv returned an invalid RGB observation: shape={image.shape}, dtype={image.dtype}") + return image + + +def language_instruction(env: Any) -> str: + instruction = str(env.get_language_instruction()) + if not instruction: + raise RuntimeError("SimplerEnv returned an empty language instruction") + return instruction + + +def close_env(env: Any) -> None: + close = getattr(env, "close", None) + if callable(close): + close()