diff --git a/.github/workflows/inference.yml b/.github/workflows/inference.yml index 10241bbfa..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,9 +30,41 @@ jobs: echo "approved=false" >> $GITHUB_OUTPUT fi - build: + runner-check: needs: gate if: needs.gate.outputs.approved == 'true' + 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: 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", } diff --git a/python/cactus/transpile/component_pipeline.py b/python/cactus/transpile/component_pipeline.py index 26eba1533..11f51cd94 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 npu_example_inputs: tuple[torch.Tensor, ...] | None = None 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 f3a277806..e15b25cce 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 @@ -58,34 +59,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