From ba95c373807fd53ace9ba2860cd81c894ac99752 Mon Sep 17 00:00:00 2001 From: Noah Cylich Date: Mon, 15 Jun 2026 14:07:39 -1000 Subject: [PATCH 1/5] convert: make gemma4_scale_factor the single source for weight reparam The gemma4 16x weight reparameterization was encoded twice: once in gemma4_scale_factor() and again as a hardcoded branch in save_tensor_with_header(). The two had to be kept in lockstep by hand, and gemma4_scale_factor() omitted the projection weights entirely. Add the vision/audio projection basenames to GEMMA4_DIV_BASENAMES and collapse the tensor_io branch to call gemma4_scale_factor(), so a single function decides every gemma4 weight's factor. Factors are identical for all converted weights (verified over 2133 names), so stored weights are byte-for-byte unchanged. Signed-off-by: Noah Cylich --- .../cactus/convert/cactus_adapters/tensor_io.py | 16 ++-------------- python/cactus/convert/model_adapters/naming.py | 2 ++ 2 files changed, 4 insertions(+), 14 deletions(-) diff --git a/python/cactus/convert/cactus_adapters/tensor_io.py b/python/cactus/convert/cactus_adapters/tensor_io.py index 8bfc632f1..cb38aa5c1 100644 --- a/python/cactus/convert/cactus_adapters/tensor_io.py +++ b/python/cactus/convert/cactus_adapters/tensor_io.py @@ -3,7 +3,7 @@ from pathlib import Path from typing import Optional, Dict, Any, List from .weight_patterns import EMBED_NAMES -from ..model_adapters.naming import GEMMA4_WEIGHT_SCALE +from ..model_adapters.naming import GEMMA4_WEIGHT_SCALE, gemma4_scale_factor try: import torch @@ -283,19 +283,7 @@ def save_tensor_with_header(tensor, output_path, precision='INT8', transpose=Fal data = data / GEMMA4_WEIGHT_SCALE if model_type == 'gemma4': - filename = output_path.name - is_audio_weight = filename.startswith('audio_') - if any(x in filename for x in ['input_norm', 'post_attn_norm', 'pre_ffn_norm', 'post_ffn_norm', - 'post_per_layer_norm', 'post_proj_norm']): - data = data / GEMMA4_WEIGHT_SCALE - elif any(x in filename for x in ['ffn_gate', 'ffn_up', 'per_layer_gate', 'moe_gate_proj', 'moe_up_proj']): - data = data * GEMMA4_WEIGHT_SCALE - elif 'router_scale' in filename: - data = data / GEMMA4_WEIGHT_SCALE - elif filename in ('token_embeddings.weights', 'output_weight.weights', 'embed_vision_proj.weights', 'embed_audio_proj.weights'): - data = data / GEMMA4_WEIGHT_SCALE - elif filename == 'output_norm.weights': - data = data * GEMMA4_WEIGHT_SCALE + data = data * gemma4_scale_factor(output_path.name) if precision in ('INT8', 'INT4'): filename = output_path.name diff --git a/python/cactus/convert/model_adapters/naming.py b/python/cactus/convert/model_adapters/naming.py index e96f568fb..b44a4a072 100644 --- a/python/cactus/convert/model_adapters/naming.py +++ b/python/cactus/convert/model_adapters/naming.py @@ -23,6 +23,8 @@ class NameMatch: GEMMA4_DIV_BASENAMES = { "token_embeddings", "output_weight", + "embed_vision_proj", + "embed_audio_proj", } From cf453c41c65adf38c122998b845332115822c8f9 Mon Sep 17 00:00:00 2001 From: Noah Cylich Date: Mon, 15 Jun 2026 14:07:39 -1000 Subject: [PATCH 2/5] transpile: reparam gemma4 projection weights for NPU export The NPU audio and vision encoders are traced directly from the in-memory torch module, which holds the original projection weights, while the soft-token 1/16 lives only in the converted embed_*_proj.weights files that the .mlpackage never reads. gemma4 multimodal embedders have no post-projection norm (only a pre-projection one), so the projection magnitude is the soft-token magnitude and nothing renormalizes it: both audio and vision soft tokens came out 16x too large on NPU. Reparameterize the projection weights with the same gemma4_scale_factor used at conversion time, applied as a reversible context manager around the NPU trace so the shared module is restored before CPU IR capture. Verified at the .mlpackage level: reparam divides both encoders' output by exactly 1/16 (cosine 1.0), matching the CPU path. Signed-off-by: Noah Cylich --- python/cactus/transpile/component_pipeline.py | 4 +- python/cactus/transpile/model_adapters.py | 27 +++++++++ python/cactus/transpile/npu/pipeline.py | 59 ++++++++++--------- 3 files changed, 61 insertions(+), 29 deletions(-) diff --git a/python/cactus/transpile/component_pipeline.py b/python/cactus/transpile/component_pipeline.py index 0f79bf5fa..6a2b3a65e 100644 --- a/python/cactus/transpile/component_pipeline.py +++ b/python/cactus/transpile/component_pipeline.py @@ -1,10 +1,11 @@ from __future__ import annotations import copy +from contextlib import AbstractContextManager from dataclasses import dataclass from dataclasses import field import re -from typing import Any +from typing import Any, Callable import numpy as np import torch @@ -30,6 +31,7 @@ class ComponentModuleSpec: metadata: dict[str, object] = field(default_factory=dict) npu_module: torch.nn.Module | None = None npu_runtime_input_count: int = 1 + npu_reparam: Callable[[torch.nn.Module], AbstractContextManager] | None = None @dataclass diff --git a/python/cactus/transpile/model_adapters.py b/python/cactus/transpile/model_adapters.py index b6e2cee0d..daac679ce 100644 --- a/python/cactus/transpile/model_adapters.py +++ b/python/cactus/transpile/model_adapters.py @@ -16,6 +16,7 @@ from cactus.transpile.component_pipeline import ComponentModuleSpec from cactus.convert.cactus_adapters.tensor_io import CACTUS_MAGIC from cactus.convert.cactus_adapters.tensor_io import align_offset +from cactus.convert.model_adapters.naming import gemma4_scale_factor _GEMMA4_SAFE_TEXT_MLP_PRODUCT_SCALE = 1.0 / 64.0 @@ -1001,6 +1002,30 @@ def _gemma4_compute_native_like_image_features( return projected +_GEMMA4_NPU_PROJ_WEIGHTS = ( + ("embed_audio", "embed_audio_proj.weights"), + ("embed_vision", "embed_vision_proj.weights"), +) + + +@contextmanager +def _gemma4_npu_projection_reparam(module: torch.nn.Module): + backbone = getattr(module, "multimodal_backbone", None) + scaled: list[tuple[torch.nn.Linear, float]] = [] + for attr, cactus_name in _GEMMA4_NPU_PROJ_WEIGHTS: + projection = getattr(getattr(backbone, attr, None), "embedding_projection", None) + factor = gemma4_scale_factor(cactus_name) + if isinstance(projection, torch.nn.Linear) and factor != 1.0: + scaled.append((projection, factor)) + for projection, factor in scaled: + projection.weight.data.mul_(factor) + try: + yield + finally: + for projection, factor in scaled: + projection.weight.data.div_(factor) + + def _gemma4_compute_native_like_audio_features( multimodal_backbone: torch.nn.Module, input_features: torch.Tensor, @@ -3745,6 +3770,7 @@ def _require(component: str) -> None: metadata={"family": "gemma4", "task": "multimodal_causal_lm_logits"}, npu_module=vision_encoder_npu, npu_runtime_input_count=2, + npu_reparam=_gemma4_npu_projection_reparam, )) if "audio_encoder" in expanded_components: specs.append(ComponentModuleSpec( @@ -3755,6 +3781,7 @@ def _require(component: str) -> None: output_keys=("audio_features",), graph_meta={**common_graph_meta, "component": "audio_encoder"}, metadata={"family": "gemma4", "task": "multimodal_causal_lm_logits"}, + npu_reparam=_gemma4_npu_projection_reparam, )) if "lm_encoder" in expanded_components: if image_features is None or audio_features is None: diff --git a/python/cactus/transpile/npu/pipeline.py b/python/cactus/transpile/npu/pipeline.py index e04429298..5f1047c9a 100644 --- a/python/cactus/transpile/npu/pipeline.py +++ b/python/cactus/transpile/npu/pipeline.py @@ -1,5 +1,6 @@ from __future__ import annotations +from contextlib import nullcontext from pathlib import Path from .audio import emit_audio_encoder_mlpackage @@ -56,34 +57,36 @@ def _resolve(per_component: int | None, default: int | None) -> int | None: qbits = component_quants[component] qdesc = f"int{qbits}" if qbits else "fp16" print(f"npu.pipeline: emitting {component} quant={qdesc}") - if component == "source_encoder": - emitted = emit_fn( - module, - bundle_root, - example_inputs=example_inputs, - input_names=tuple(getattr(spec, "input_keys", ()) or ()), - filename=filename, - quantize_bits=qbits, - ) - elif component == "vision_encoder": - names = ("x",) + tuple(spec.input_keys[1:n_runtime]) - emitted = emit_fn( - module, - bundle_root, - runtime_inputs=tuple(zip(names, example_inputs[:n_runtime])), - baked_inputs=example_inputs[n_runtime:], - filename=filename, - quantize_bits=qbits, - ) - else: - emitted = emit_fn( - module, - bundle_root, - example_input=example_inputs[0], - baked_inputs=example_inputs[1:], - filename=filename, - quantize_bits=qbits, - ) + reparam = getattr(spec, "npu_reparam", None) + with (reparam(module) if reparam else nullcontext()): + if component == "source_encoder": + emitted = emit_fn( + module, + bundle_root, + example_inputs=example_inputs, + input_names=tuple(getattr(spec, "input_keys", ()) or ()), + filename=filename, + quantize_bits=qbits, + ) + elif component == "vision_encoder": + names = ("x",) + tuple(spec.input_keys[1:n_runtime]) + emitted = emit_fn( + module, + bundle_root, + runtime_inputs=tuple(zip(names, example_inputs[:n_runtime])), + baked_inputs=example_inputs[n_runtime:], + filename=filename, + quantize_bits=qbits, + ) + else: + emitted = emit_fn( + module, + bundle_root, + example_input=example_inputs[0], + baked_inputs=example_inputs[1:], + filename=filename, + quantize_bits=qbits, + ) if emitted: results[f"npu_{component}"] = f"components/{emitted}" return results From 713374e8ecb4cbd53a52ed0b98f08905e7781090 Mon Sep 17 00:00:00 2001 From: Noah Cylich Date: Tue, 16 Jun 2026 10:02:44 -1000 Subject: [PATCH 3/5] Update inference.yml Signed-off-by: Noah Cylich --- .github/workflows/inference.yml | 26 ++++++++++++++++---------- 1 file changed, 16 insertions(+), 10 deletions(-) diff --git a/.github/workflows/inference.yml b/.github/workflows/inference.yml index 10241bbfa..281e6dd01 100644 --- a/.github/workflows/inference.yml +++ b/.github/workflows/inference.yml @@ -32,18 +32,13 @@ jobs: build: needs: gate if: needs.gate.outputs.approved == 'true' - runs-on: [self-hosted, linux, arm64] + runs-on: ubuntu-24.04-arm timeout-minutes: 60 steps: - uses: actions/checkout@v4 - with: - clean: false - - name: Setup Python environment - run: | - if [ ! -f venv/bin/activate ]; then - source ./setup - fi + - name: Install build tools + run: sudo apt-get update && sudo apt-get install -y cmake build-essential libcurl4-openssl-dev - name: Build libcactus + engine tests run: | @@ -55,7 +50,7 @@ jobs: model-tests: needs: build - runs-on: [self-hosted, linux, arm64] + runs-on: ubuntu-24.04-arm timeout-minutes: 60 env: CACTUS_NO_CLOUD_TELE: "1" @@ -67,8 +62,19 @@ jobs: - openai/whisper-small steps: - uses: actions/checkout@v4 + + - name: Install build tools + run: sudo apt-get update && sudo apt-get install -y cmake build-essential libcurl4-openssl-dev + + - name: Build libcactus + run: bash cactus-engine/build.sh + + - uses: actions/setup-python@v5 with: - clean: false + python-version: "3.12" + + - name: Install dependencies + run: pip install -e "python/[dev,serve]" - name: Run tests for ${{ matrix.model }} env: From 37aa0a4e7cbade654b8035b08e3c90d5c79a5af2 Mon Sep 17 00:00:00 2001 From: Noah Cylich Date: Tue, 16 Jun 2026 10:27:47 -1000 Subject: [PATCH 4/5] ci: create venv for inference tests Signed-off-by: Noah Cylich --- .github/workflows/inference.yml | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/.github/workflows/inference.yml b/.github/workflows/inference.yml index 281e6dd01..4abbcab33 100644 --- a/.github/workflows/inference.yml +++ b/.github/workflows/inference.yml @@ -73,8 +73,12 @@ jobs: with: python-version: "3.12" - - name: Install dependencies - run: pip install -e "python/[dev,serve]" + - name: Setup Python environment + run: | + python -m venv venv + source venv/bin/activate + python -m pip install --upgrade pip + pip install -e "python/[dev,serve]" - name: Run tests for ${{ matrix.model }} env: From f4fb02e0c6fdcebf5a08b85b180b6a350ef837d9 Mon Sep 17 00:00:00 2001 From: Noah Cylich Date: Tue, 16 Jun 2026 12:14:17 -1000 Subject: [PATCH 5/5] ci: skip inference tests when self-hosted runner is unavailable Signed-off-by: Noah Cylich --- .github/workflows/inference.yml | 65 ++++++++++++++++++++++----------- 1 file changed, 44 insertions(+), 21 deletions(-) diff --git a/.github/workflows/inference.yml b/.github/workflows/inference.yml index 4abbcab33..d8cac3705 100644 --- a/.github/workflows/inference.yml +++ b/.github/workflows/inference.yml @@ -7,6 +7,7 @@ on: permissions: contents: read + actions: read concurrency: group: inference-${{ github.event.pull_request.number || github.ref }} @@ -29,16 +30,53 @@ jobs: echo "approved=false" >> $GITHUB_OUTPUT fi - build: + runner-check: needs: gate if: needs.gate.outputs.approved == 'true' - runs-on: ubuntu-24.04-arm + runs-on: ubuntu-latest + timeout-minutes: 5 + outputs: + available: ${{ steps.check.outputs.available }} + steps: + - id: check + env: + GH_TOKEN: ${{ github.token }} + REPO: ${{ github.repository }} + run: | + set -euo pipefail + status=$(curl -sS -L -w "%{http_code}" -o runners.json \ + -H "Authorization: Bearer ${GH_TOKEN}" \ + -H "Accept: application/vnd.github+json" \ + -H "X-GitHub-Api-Version: 2022-11-28" \ + "https://api.github.com/repos/${REPO}/actions/runners?per_page=100" || true) + if [ "$status" != "200" ]; then + echo "available=false" >> "$GITHUB_OUTPUT" + echo "::notice::Could not inspect self-hosted runners (HTTP $status); skipping inference tests instead of queueing." + exit 0 + fi + count=$(jq '[.runners[] | select(.status == "online" and .busy == false) | select(([.labels[].name] | contains(["self-hosted", "linux", "arm64"])))] | length' runners.json) + if [ "$count" -gt 0 ]; then + echo "available=true" >> "$GITHUB_OUTPUT" + else + echo "available=false" >> "$GITHUB_OUTPUT" + echo "::notice::No idle self-hosted linux arm64 runner is available; skipping inference tests." + fi + + build: + needs: [gate, runner-check] + if: needs.gate.outputs.approved == 'true' && needs.runner-check.outputs.available == 'true' + runs-on: [self-hosted, linux, arm64] timeout-minutes: 60 steps: - uses: actions/checkout@v4 + with: + clean: false - - name: Install build tools - run: sudo apt-get update && sudo apt-get install -y cmake build-essential libcurl4-openssl-dev + - name: Setup Python environment + run: | + if [ ! -f venv/bin/activate ]; then + source ./setup + fi - name: Build libcactus + engine tests run: | @@ -50,7 +88,7 @@ jobs: model-tests: needs: build - runs-on: ubuntu-24.04-arm + runs-on: [self-hosted, linux, arm64] timeout-minutes: 60 env: CACTUS_NO_CLOUD_TELE: "1" @@ -62,23 +100,8 @@ jobs: - openai/whisper-small steps: - uses: actions/checkout@v4 - - - name: Install build tools - run: sudo apt-get update && sudo apt-get install -y cmake build-essential libcurl4-openssl-dev - - - name: Build libcactus - run: bash cactus-engine/build.sh - - - uses: actions/setup-python@v5 with: - python-version: "3.12" - - - name: Setup Python environment - run: | - python -m venv venv - source venv/bin/activate - python -m pip install --upgrade pip - pip install -e "python/[dev,serve]" + clean: false - name: Run tests for ${{ matrix.model }} env: