diff --git a/docs/model-catalog.md b/docs/model-catalog.md index 662139929..8339473b9 100644 --- a/docs/model-catalog.md +++ b/docs/model-catalog.md @@ -197,7 +197,7 @@ Also registered with `T5ForConditionalGeneration` (task: `seq2seq`): | `qwen3_tts` | `Qwen3TTSForConditionalGeneration` | `tts` | — | | `qwen3_tts_tokenizer_12hz` | `Qwen3TTSTokenizerV2Model` | `codec` | — | | `vibevoice` | `VibeVoiceForConditionalGeneration` | `vibevoice-tts` | `vibevoice/VibeVoice-1.5B-hf` | -| `VibeVoiceForASRTraining` | `VibeVoiceASRForConditionalGeneration` | `vibevoice-asr` | [`microsoft/VibeVoice-ASR`](vibevoice-asr.md) | +| `VibeVoiceForASRTraining` | `VibeVoiceASRForConditionalGeneration` | `vibevoice-asr` | [`microsoft/VibeVoice-ASR`](vibevoice-asr.md); [`microsoft/VibeVoice-ASR-BitNet`](vibevoice-asr.md) (pinned dense-F32 conversion source; native I2_S/I8_S GGUF rejected) | ### Audio Feature Extraction diff --git a/docs/vibevoice-asr.md b/docs/vibevoice-asr.md index fcad29f36..c65cf4d73 100644 --- a/docs/vibevoice-asr.md +++ b/docs/vibevoice-asr.md @@ -8,6 +8,22 @@ the shared `vibevoice` configuration declares `VibeVoiceForConditionalGeneration`; unknown or streaming VibeVoice architectures fail closed rather than falling through to either implementation. +## BitNet conversion source + +[`microsoft/VibeVoice-ASR-BitNet`](https://huggingface.co/microsoft/VibeVoice-ASR-BitNet) +is supported only at revision `66e78021ab8f5f06133d1ab421ba4d348bda97c9`. It uses the +same staged ASR package and host contract, but it is a distinct artifact format: +Mobius streams its three dense F32 safetensors shards as the ONNX conversion source +and records that provenance in `weight-loading-report.json`. + +The release's `vibeasr-lm-i2_s-embed-q6_k.gguf` and +`vibeasr-vae-encoder-i8_s.gguf` files are **not** supported import inputs. Their +I2_S and I8_S tensor types require VibeASR.cpp-specific packed storage and custom +CPU kernels; converting them through generic GGUF or affine ONNX quantization would +not preserve their execution semantics. Mobius identifies these exact native +artifacts during local import and Hub preflight, then fails with an explicit error +rather than claiming a native BitNet or dequantized execution path. + This is distinct from the streaming-ASR work tracked in #723. It uses the official offline model's Qwen2 decoder (3584 hidden width, 28 layers, 28 query heads, 4 KV heads) and its 24 kHz, 3200-sample waveform framing. diff --git a/src/mobius/integrations/_vibeasr_bitnet.py b/src/mobius/integrations/_vibeasr_bitnet.py new file mode 100644 index 000000000..81344fb15 --- /dev/null +++ b/src/mobius/integrations/_vibeasr_bitnet.py @@ -0,0 +1,345 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Artifact identities and import verdicts for VibeVoice ASR BitNet. + +The official VibeVoice ASR BitNet release contains an F32 safetensors +conversion source alongside two execution-only GGUF files. The latter are not +ordinary llama.cpp quantizations: ``I2_S`` needs VibeASR.cpp's packed ternary +and activation kernels, and ``I8_S`` needs its fused VAE kernels. This module +centralizes that artifact-level distinction without teaching generic ONNX +components about a model-specific execution format. The dense safetensors +conversion source is recorded separately so its ONNX path cannot be confused +with VibeASR.cpp-native quantized execution. +""" + +from __future__ import annotations + +__all__ = [ + "VIBEVOICE_ASR_BITNET_DENSE_F32_TENSOR_COUNT", + "VIBEVOICE_ASR_BITNET_DENSE_F32_VALUE_COUNT", + "VIBEVOICE_ASR_BITNET_DENSE_SAFETENSORS", + "VIBEVOICE_ASR_BITNET_ARTIFACTS", + "VIBEVOICE_ASR_BITNET_REPOSITORY", + "VIBEVOICE_ASR_BITNET_REVISION", + "build_vibeasr_bitnet_dense_weight_plan", + "is_vibeasr_bitnet_conversion_source", + "VibeASRBitNetGGUFArtifact", + "VibeASRBitNetSafetensorsArtifact", + "find_vibeasr_bitnet_gguf_artifact", + "reject_vibeasr_bitnet_gguf", +] + +import math +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import PurePath +from typing import TYPE_CHECKING + +from mobius.integrations.gguf._errors import VibeASRBitNetGGUFImportError +from mobius.integrations.gguf._header import GGUFHeaderInfo + +if TYPE_CHECKING: + import onnx_ir as ir + + from mobius.integrations._weight_loading import StreamingWeightPlan + from mobius.models.vibevoice_asr import VibeVoiceASRForConditionalGeneration + +VIBEVOICE_ASR_BITNET_REPOSITORY = "microsoft/VibeVoice-ASR-BitNet" +VIBEVOICE_ASR_BITNET_REVISION = "66e78021ab8f5f06133d1ab421ba4d348bda97c9" +VIBEVOICE_ASR_BITNET_DENSE_F32_TENSOR_COUNT = 1_177 +VIBEVOICE_ASR_BITNET_DENSE_F32_VALUE_COUNT = 2_814_116_321 +_VIBEVOICE_ASR_BITNET_INDEX_REPORTED_PARAMETER_COUNT = 322_592_829 +_VIBEASR_LM_FILE_TYPE = 40 +_VIBEASR_VAE_FILE_TYPE = 41 +_VIBEASR_I2_S_TYPE_ID = 36 +_VIBEASR_I8_S_TYPE_ID = 37 + + +@dataclass(frozen=True, slots=True) +class VibeASRBitNetSafetensorsArtifact: + """Immutable identity of one dense F32 conversion-source shard.""" + + filename: str + size_bytes: int + sha256: str + + +@dataclass(frozen=True, slots=True) +class VibeASRBitNetGGUFArtifact: + """Pinned fingerprint and native execution contract of one GGUF artifact.""" + + filename: str + role: str + architecture: str + tensor_count: int + size_bytes: int + sha256: str + native_format: str + blocker: str + file_type: int + tensor_type_ids: frozenset[int] + + +VIBEVOICE_ASR_BITNET_DENSE_SAFETENSORS = ( + VibeASRBitNetSafetensorsArtifact( + filename="model-00001-of-00003.safetensors", + size_bytes=4_996_674_400, + sha256="58cb328634bb4b7e5afcc4f14c43261a0c636c9031b0b085bd3ef53c131aaf19", + ), + VibeASRBitNetSafetensorsArtifact( + filename="model-00002-of-00003.safetensors", + size_bytes=4_963_184_724, + sha256="410349401256a9a995424408f8f77a21cad5e53183e07207820d7b25e84682df", + ), + VibeASRBitNetSafetensorsArtifact( + filename="model-00003-of-00003.safetensors", + size_bytes=1_296_761_656, + sha256="49c2d7591f9dcbbd6fb37baa58379211f21df30dbd3c6ff39cb8c81473660cdb", + ), +) + + +_LM_BLOCKER = ( + "The I2_S language-model projections use VibeASR.cpp's packed ternary code " + "layout, tensor-sidecar scales, and ISA-specific I2_S-by-I8_S activation kernel. " + "ORT MatMulNBits is affine and cannot preserve or execute this contract." +) +_VAE_BLOCKER = ( + "The I8_S acoustic/semantic encoder and connector weights rely on VibeASR.cpp's " + "all-INT8 fused convolution, MatMul, RMSNorm, and residual kernels. There is no " + "lossless ONNX Runtime import or execution provider for that contract." +) + +VIBEVOICE_ASR_BITNET_ARTIFACTS = ( + VibeASRBitNetGGUFArtifact( + filename="vibeasr-lm-i2_s-embed-q6_k.gguf", + role="decoder", + architecture="qwen2", + tensor_count=339, + size_bytes=992_877_600, + sha256="fbe273d8dc2f2433bb25f849e19d77ea65aaa2188d12c20cee987ab6f321e002", + native_format="I2_S ternary projections with Q6_K embedding/head", + blocker=_LM_BLOCKER, + file_type=_VIBEASR_LM_FILE_TYPE, + tensor_type_ids=frozenset({0, 1, 14, _VIBEASR_I2_S_TYPE_ID}), + ), + VibeASRBitNetGGUFArtifact( + filename="vibeasr-vae-encoder-i8_s.gguf", + role="acoustic, semantic, and connector stages", + architecture="vibeasr-vae", + tensor_count=562, + size_bytes=703_080_064, + sha256="4941c82608c253ec066b5cc74d3dd11a5c8fef96cccbc5b87359ef0fe4338df6", + native_format="I8_S fused VAE encoder and connector kernels", + blocker=_VAE_BLOCKER, + file_type=_VIBEASR_VAE_FILE_TYPE, + tensor_type_ids=frozenset({0, _VIBEASR_I8_S_TYPE_ID}), + ), +) + + +def is_vibeasr_bitnet_conversion_source(model_id: str, revision: str | None) -> bool: + """Return whether the requested immutable source is the audited F32 release.""" + return ( + model_id == VIBEVOICE_ASR_BITNET_REPOSITORY + and revision == VIBEVOICE_ASR_BITNET_REVISION + ) + + +def build_vibeasr_bitnet_dense_weight_plan( + model: VibeVoiceASRForConditionalGeneration, + source_tensors: Mapping[str, tuple[str, list[int], str]], + _initializers: Mapping[str, ir.Value], +) -> StreamingWeightPlan: + """Classify every source tensor for the staged dense-F32 conversion route. + + The parent ASR module remains authoritative for HF-to-ONNX name alignment. + Marker tensors exercise that mapping without materializing any checkpoint + values; every non-decoder source must map to an exported initializer. + """ + import torch + + from mobius.integrations._weight_loading import StreamingWeightPlan, StreamingWeightSource + + if len(source_tensors) != VIBEVOICE_ASR_BITNET_DENSE_F32_TENSOR_COUNT: + raise ValueError( + "VibeVoice ASR BitNet dense source tensor count changed: expected " + f"{VIBEVOICE_ASR_BITNET_DENSE_F32_TENSOR_COUNT}, got {len(source_tensors)}." + ) + non_f32 = { + source_name: source_dtype + for source_name, (_, _, source_dtype) in source_tensors.items() + if source_dtype != "F32" + } + if non_f32: + examples = sorted(non_f32.items())[:5] + raise ValueError( + "VibeVoice ASR BitNet dense conversion source must contain only F32 tensors; " + f"found {examples}." + ) + value_count = sum(math.prod(shape) for _, shape, _ in source_tensors.values()) + if value_count != VIBEVOICE_ASR_BITNET_DENSE_F32_VALUE_COUNT: + raise ValueError( + "VibeVoice ASR BitNet dense source value count changed: expected " + f"{VIBEVOICE_ASR_BITNET_DENSE_F32_VALUE_COUNT}, got {value_count}." + ) + + markers = {name: torch.empty(0) for name in source_tensors} + mapped = model.preprocess_weights(markers) + marker_sources = {id(marker): name for name, marker in markers.items()} + targets: dict[str, StreamingWeightSource] = {} + for target_name, marker in mapped.items(): + source_name = marker_sources.get(id(marker)) + if source_name is None: + raise ValueError( + "VibeVoice ASR BitNet weight preprocessing transformed a source marker for " + f"{target_name!r}; streaming requires a one-to-one source tensor mapping." + ) + if target_name in targets: + raise ValueError( + f"VibeVoice ASR BitNet maps multiple source tensors to {target_name!r}." + ) + targets[target_name] = StreamingWeightSource( + source_name=source_name, + expected_dtype="F32", + ) + + used_sources = {source.source_name for source in targets.values()} + ignored: dict[str, str] = {} + for source_name in source_tensors: + if source_name in used_sources: + continue + if source_name.startswith("model.acoustic_tokenizer.decoder."): + ignored[source_name] = ( + "The source acoustic VAE decoder is not an ASR inference stage." + ) + continue + raise ValueError( + f"VibeVoice ASR BitNet source tensor {source_name!r} is not classified by " + "the staged inference package." + ) + + return StreamingWeightPlan( + targets=targets, + ignored=ignored, + report={ + "source_format": "safetensors", + "source_storage_dtype": "float32", + "source_tensor_count": len(source_tensors), + "source_value_count": value_count, + "source_parameter_count_status": ( + "The checkpoint index's reported 322592829 parameter count is inconsistent " + "with the exact dense-F32 tensor-byte census (2814116321 values); the latter " + "is authoritative for this export." + ), + "index_reported_parameter_count": _VIBEVOICE_ASR_BITNET_INDEX_REPORTED_PARAMETER_COUNT, + "native_bitnet_execution": False, + "native_gguf_disposition": ( + "not imported; VibeASR.cpp I2_S/I8_S kernels and packing are unsupported" + ), + "native_gguf_artifacts": [ + { + "filename": artifact.filename, + "sha256": artifact.sha256, + "disposition": "unsupported_native_execution", + } + for artifact in VIBEVOICE_ASR_BITNET_ARTIFACTS + ], + }, + ) + + +def _is_vibeasr_lm_header(header: GGUFHeaderInfo) -> bool: + """Identify the custom LM by its model-owned header identity, not Qwen2 alone.""" + return ( + header.architecture == "qwen2" + and header.file_type == _VIBEASR_LM_FILE_TYPE + and _VIBEASR_I2_S_TYPE_ID in header.tensor_type_ids + ) + + +def _is_vibeasr_vae_header(header: GGUFHeaderInfo) -> bool: + """Identify the VAE only when its custom all-INT8 storage is present.""" + return ( + header.architecture == "vibeasr-vae" + and header.file_type == _VIBEASR_VAE_FILE_TYPE + and _VIBEASR_I8_S_TYPE_ID in header.tensor_type_ids + ) + + +def find_vibeasr_bitnet_gguf_artifact( + *, + repository: str | None = None, + revision: str | None = None, + filename: str | None = None, + size_bytes: int | None = None, + sha256: str | None = None, + header: GGUFHeaderInfo | None = None, +) -> VibeASRBitNetGGUFArtifact | None: + """Return a verified VibeASR native artifact, never a generic Qwen2 alias. + + An exact Hub artifact is matched by repository, immutable revision, basename, + and optional size/checksum. A local file is recognized only by VibeASR-owned + header identity; the generic ``qwen2`` architecture is intentionally + insufficient because ordinary Qwen2 GGUF files remain supported. + """ + basename = PurePath(filename).name if filename else None + if ( + repository == VIBEVOICE_ASR_BITNET_REPOSITORY + and revision == VIBEVOICE_ASR_BITNET_REVISION + and basename is not None + ): + for artifact in VIBEVOICE_ASR_BITNET_ARTIFACTS: + if basename != artifact.filename: + continue + if size_bytes is not None and size_bytes != artifact.size_bytes: + continue + if sha256 is not None and sha256.casefold() != artifact.sha256: + continue + return artifact + if header is not None: + for artifact in VIBEVOICE_ASR_BITNET_ARTIFACTS: + is_match = ( + header.architecture == artifact.architecture + and header.file_type == artifact.file_type + and ( + _is_vibeasr_vae_header(header) + if artifact.architecture == "vibeasr-vae" + else _is_vibeasr_lm_header(header) + ) + ) + if is_match: + return artifact + return None + + +def reject_vibeasr_bitnet_gguf( + *, + source: str, + repository: str | None = None, + revision: str | None = None, + filename: str | None = None, + size_bytes: int | None = None, + sha256: str | None = None, + header: GGUFHeaderInfo | None = None, +) -> None: + """Fail closed before config extraction or tensor payload access.""" + artifact = find_vibeasr_bitnet_gguf_artifact( + repository=repository, + revision=revision, + filename=filename, + size_bytes=size_bytes, + sha256=sha256, + header=header, + ) + if artifact is None: + return + raise VibeASRBitNetGGUFImportError( + f"Direct GGUF import is unsupported for VibeVoice ASR BitNet {artifact.role} " + f"artifact {artifact.filename!r} ({artifact.native_format}) from {source!r}. " + f"{artifact.blocker} Build the pinned Hugging Face safetensors checkpoint with " + "`mobius build` instead; that route uses the release's dense F32 conversion " + "source and does not claim native BitNet/GGUF preservation or execution. " + "No ONNX artifacts were emitted." + ) diff --git a/src/mobius/integrations/_vibeasr_bitnet_test.py b/src/mobius/integrations/_vibeasr_bitnet_test.py new file mode 100644 index 000000000..faa05ce5f --- /dev/null +++ b/src/mobius/integrations/_vibeasr_bitnet_test.py @@ -0,0 +1,403 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Tests for the fail-closed VibeVoice ASR BitNet native-GGUF verdict.""" + +from __future__ import annotations + +import json +import struct +from pathlib import Path +from types import SimpleNamespace +from unittest import mock + +import pytest + +from mobius.integrations._vibeasr_bitnet import ( + VIBEVOICE_ASR_BITNET_ARTIFACTS, + VIBEVOICE_ASR_BITNET_DENSE_F32_TENSOR_COUNT, + VIBEVOICE_ASR_BITNET_DENSE_F32_VALUE_COUNT, + VIBEVOICE_ASR_BITNET_DENSE_SAFETENSORS, + VIBEVOICE_ASR_BITNET_REPOSITORY, + VIBEVOICE_ASR_BITNET_REVISION, + build_vibeasr_bitnet_dense_weight_plan, + find_vibeasr_bitnet_gguf_artifact, +) +from mobius.integrations.gguf._errors import VibeASRBitNetGGUFImportError +from mobius.integrations.gguf._header import ( + GGUFHeaderInfo, + _gguf_header_info_from_header, +) +from mobius.integrations.gguf._reader import GGUFModel + + +def _string(value: str) -> bytes: + encoded = value.encode("utf-8") + return struct.pack(" bytes: + return _string(key) + struct.pack(" bytes: + return _string(key) + struct.pack(" bytes: + # The import verdict needs only type IDs, so zero-length tensor descriptors + # exercise the bounded header path without requiring tensor payloads. + return _string(name) + struct.pack(" bytes: + metadata = ( + _string_entry("general.architecture", artifact.architecture), + _string_entry( + "general.name", + "models" if artifact.architecture == "qwen2" else "VibeASR VAE Encoder", + ), + _uint32_entry("general.file_type", artifact.file_type), + _uint32_entry("general.quantization_version", 2), + ) + tensors = [ + _tensor_entry(f"tensor-{index}", type_id) + for index, type_id in enumerate(sorted(artifact.tensor_type_ids)) + ] + return ( + b"GGUF" + + struct.pack(" None: + header = _gguf_header_info_from_header( + _native_header(artifact), + source=artifact.filename, + collect_tensor_type_ids=True, + ) + + assert header.architecture == artifact.architecture + assert header.file_type == artifact.file_type + assert header.quantization_version == 2 + assert header.tensor_type_ids == artifact.tensor_type_ids + assert find_vibeasr_bitnet_gguf_artifact(header=header) == artifact + + +@pytest.mark.parametrize("artifact", VIBEVOICE_ASR_BITNET_ARTIFACTS) +def test_pinned_hub_artifact_fingerprint_is_exact(artifact) -> None: + assert ( + find_vibeasr_bitnet_gguf_artifact( + repository=VIBEVOICE_ASR_BITNET_REPOSITORY, + revision=VIBEVOICE_ASR_BITNET_REVISION, + filename=artifact.filename, + size_bytes=artifact.size_bytes, + sha256=artifact.sha256, + ) + == artifact + ) + assert ( + find_vibeasr_bitnet_gguf_artifact( + repository=VIBEVOICE_ASR_BITNET_REPOSITORY, + revision=VIBEVOICE_ASR_BITNET_REVISION, + filename=artifact.filename, + size_bytes=artifact.size_bytes + 1, + sha256=artifact.sha256, + ) + is None + ) + + +@pytest.mark.parametrize("artifact", VIBEVOICE_ASR_BITNET_ARTIFACTS) +def test_local_reader_rejects_native_gguf_before_reader_payload_access( + tmp_path: Path, artifact +) -> None: + path = tmp_path / artifact.filename + path.write_bytes(_native_header(artifact)) + + with pytest.raises( + VibeASRBitNetGGUFImportError, + match=r"Direct GGUF import is unsupported.*No ONNX artifacts were emitted", + ): + GGUFModel(path) + + +@pytest.mark.parametrize("artifact", VIBEVOICE_ASR_BITNET_ARTIFACTS) +def test_hub_header_preflight_rejects_before_download(artifact) -> None: + from mobius.integrations.gguf import _builder + + response = mock.MagicMock() + response.iter_bytes.return_value = [_native_header(artifact)] + response_context = mock.MagicMock() + response_context.__enter__.return_value = response + session = mock.MagicMock() + session.stream.return_value = response_context + + with ( + mock.patch.object( + _builder, + "get_hf_file_metadata", + return_value=SimpleNamespace( + commit_hash=VIBEVOICE_ASR_BITNET_REVISION, + location="https://cdn.example/native.gguf", + ), + ), + mock.patch.object(_builder, "get_session", return_value=session), + pytest.raises(VibeASRBitNetGGUFImportError, match=artifact.filename), + ): + _builder._preflight_hf_gguf_file( + VIBEVOICE_ASR_BITNET_REPOSITORY, + artifact.filename, + revision=VIBEVOICE_ASR_BITNET_REVISION, + ) + + +@pytest.mark.parametrize("artifact", VIBEVOICE_ASR_BITNET_ARTIFACTS) +def test_builder_validation_rejects_native_profile_before_config_extraction(artifact) -> None: + from mobius.integrations.gguf._builder import _validate_gguf_model + + class NativeGGUF: + def __init__(self) -> None: + self.architecture = artifact.architecture + self.tensor_names = ["native-projection"] + + def get_metadata(self, key: str, default=None): + return { + "general.file_type": artifact.file_type, + "general.quantization_version": 2, + }.get(key, default) + + def reader_tensors(self): + return [ + SimpleNamespace(tensor_type=SimpleNamespace(value=type_id)) + for type_id in artifact.tensor_type_ids + ] + + with pytest.raises(VibeASRBitNetGGUFImportError, match=artifact.native_format): + _validate_gguf_model(NativeGGUF(), source=artifact.filename) + + +@pytest.mark.parametrize("requested_revision", [None, VIBEVOICE_ASR_BITNET_REVISION]) +def test_hub_preflight_reports_both_native_artifacts_without_download( + monkeypatch: pytest.MonkeyPatch, requested_revision: str | None +) -> None: + from mobius.integrations.gguf._preflight import preflight_hf_gguf + + artifacts = {artifact.filename: artifact for artifact in VIBEVOICE_ASR_BITNET_ARTIFACTS} + + class MetadataOnlyApi: + def list_repo_files(self, repo_id, revision=None, token=None): + assert repo_id == VIBEVOICE_ASR_BITNET_REPOSITORY + assert revision == requested_revision + return list(artifacts) + + def get_paths_info(self, repo_id, paths, revision=None, token=None, expand=False): + assert repo_id == VIBEVOICE_ASR_BITNET_REPOSITORY + assert revision == requested_revision + assert expand + return [ + SimpleNamespace( + path=filename, + size=artifacts[filename].size_bytes, + lfs=SimpleNamespace(sha256=artifacts[filename].sha256), + ) + for filename in paths + ] + + def model_info(self, repo_id, revision=None, token=None, expand=None): + assert repo_id == VIBEVOICE_ASR_BITNET_REPOSITORY + assert revision == requested_revision + assert expand == ["gguf", "sha"] + return SimpleNamespace( + sha=VIBEVOICE_ASR_BITNET_REVISION, + gguf={"architecture": "qwen2", "total": 1_777_088_000}, + ) + + import huggingface_hub + + monkeypatch.setattr(huggingface_hub, "HfApi", lambda *args, **kwargs: MetadataOnlyApi()) + report = preflight_hf_gguf( + VIBEVOICE_ASR_BITNET_REPOSITORY, + revision=requested_revision, + ) + + assert report.total_tensors is None + assert report.total_params == 1_777_088_000 + assert [(file.filename, file.size_bytes, file.sha256) for file in report.files] == [ + (artifact.filename, artifact.size_bytes, artifact.sha256) + for artifact in VIBEVOICE_ASR_BITNET_ARTIFACTS + ] + assert len(report.blockers) == len(VIBEVOICE_ASR_BITNET_ARTIFACTS) + assert all( + artifact.filename in " ".join(report.blockers) for artifact in artifacts.values() + ) + assert not report.exportable + + +def test_generic_qwen2_with_a_standard_q1_header_is_not_a_vibeasr_alias() -> None: + header = GGUFHeaderInfo( + architecture="qwen2", + tensor_count=1, + split_no=None, + split_count=None, + split_tensors_count=None, + file_type=40, + tensor_type_ids=frozenset({0, 41}), + ) + + assert find_vibeasr_bitnet_gguf_artifact(header=header) is None + + +@pytest.mark.arch_validation +def test_pinned_dense_f32_index_classifies_all_asr_source_tensors() -> None: + """The public pinned index and the five-stage graph agree on every weight role.""" + from huggingface_hub import hf_hub_download + + from mobius.integrations.transformers._builder import build_transformers_model + from mobius.models import VibeVoiceASRForConditionalGeneration + + index_path = hf_hub_download( + VIBEVOICE_ASR_BITNET_REPOSITORY, + "model.safetensors.index.json", + revision=VIBEVOICE_ASR_BITNET_REVISION, + ) + with Path(index_path).open(encoding="utf-8") as file: + weight_map = json.load(file)["weight_map"] + assert len(weight_map) == VIBEVOICE_ASR_BITNET_DENSE_F32_TENSOR_COUNT + assert set(weight_map.values()) == { + artifact.filename for artifact in VIBEVOICE_ASR_BITNET_DENSE_SAFETENSORS + } + + package = build_transformers_model( + VIBEVOICE_ASR_BITNET_REPOSITORY, + revision=VIBEVOICE_ASR_BITNET_REVISION, + load_weights=False, + ) + initializers = { + name: initializer + for model in package.values() + for name, initializer in model.graph.initializers.items() + if initializer.const_value is None + } + source_tensors = { + source_name: (weight_map[source_name], [1], "F32") for source_name in weight_map + } + first_source = next(iter(source_tensors)) + source_tensors[first_source] = ( + weight_map[first_source], + [VIBEVOICE_ASR_BITNET_DENSE_F32_VALUE_COUNT - len(source_tensors) + 1], + "F32", + ) + plan = build_vibeasr_bitnet_dense_weight_plan( + VibeVoiceASRForConditionalGeneration(package.config), + source_tensors, + initializers, + ) + + assert len(plan.targets) == 901 + assert len(plan.ignored) == 276 + assert set(plan.targets).issubset(initializers) + assert all( + source.expected_dtype == "F32" and source.mode == "direct" + for source in plan.targets.values() + ) + assert all(name.startswith("model.acoustic_tokenizer.decoder.") for name in plan.ignored) + assert plan.report["source_value_count"] == VIBEVOICE_ASR_BITNET_DENSE_F32_VALUE_COUNT + assert plan.report["native_bitnet_execution"] is False + + +def test_builder_pins_and_selects_dense_streaming_route( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import onnx_ir as ir + + from mobius._model_package import ModelPackage + from mobius._testing import make_config + from mobius.integrations import _vibeasr_bitnet + from mobius.integrations.transformers import _builder, _config_resolver + + parent_config = SimpleNamespace(model_type="qwen2") + config = make_config(model_type="qwen2") + model = ir.Model( + ir.Graph([], [], nodes=[], name="model"), + ir_version=11, + ) + package = ModelPackage({"decoder": model}) + config_calls = [] + streaming_calls = [] + + class Module: + def __init__(self, config) -> None: + self.config = config + + def load_config(model_id, **kwargs): + config_calls.append((model_id, kwargs)) + return parent_config, False + + monkeypatch.setattr(_builder, "_load_transformers_config", load_config) + monkeypatch.setattr( + _builder, + "_select_primary_config", + lambda value: (value, value, "qwen2"), + ) + monkeypatch.setattr( + _builder, + "_resolve_module_class", + lambda *args, **kwargs: (Module, "text-generation", "qwen2"), + ) + monkeypatch.setattr(_config_resolver, "_config_from_hf", lambda *args, **kwargs: config) + monkeypatch.setattr(_builder, "build_from_module", lambda *args, **kwargs: package) + monkeypatch.setattr( + _vibeasr_bitnet, + "build_vibeasr_bitnet_dense_weight_plan", + lambda *args, **kwargs: None, + ) + + report = { + "format": "mobius.weight-loading-report.v1", + "output_weight_format": "dense", + "native_fp8": False, + "native_bitnet_execution": False, + } + + def stream(*args, **kwargs): + streaming_calls.append((args, kwargs)) + return report + + monkeypatch.setattr(_builder, "stream_preprocessed_safetensors_to_package", stream) + result = _builder.build_transformers_model(VIBEVOICE_ASR_BITNET_REPOSITORY) + + assert result is package + assert config_calls == [ + ( + VIBEVOICE_ASR_BITNET_REPOSITORY, + { + "revision": VIBEVOICE_ASR_BITNET_REVISION, + "trust_remote_code": False, + }, + ) + ] + assert streaming_calls[0][0][:2] == (package, VIBEVOICE_ASR_BITNET_REPOSITORY) + assert streaming_calls[0][1]["revision"] == VIBEVOICE_ASR_BITNET_REVISION + assert package.weight_loading_report == report + assert model.metadata_props["mobius.source_revision"] == VIBEVOICE_ASR_BITNET_REVISION + + +def test_builder_rejects_an_unpinned_bitnet_revision(monkeypatch: pytest.MonkeyPatch) -> None: + from mobius.integrations.transformers import _builder + + monkeypatch.setattr( + _builder, + "_load_transformers_config", + lambda *args, **kwargs: pytest.fail("the unsupported revision must not load config"), + ) + + with pytest.raises(ValueError, match="supported only at the audited revision"): + _builder.build_transformers_model( + VIBEVOICE_ASR_BITNET_REPOSITORY, + revision="not-the-pinned-artifact", + load_weights=False, + ) diff --git a/src/mobius/integrations/_weight_loading.py b/src/mobius/integrations/_weight_loading.py index 95cd9bdb9..4d16aa04c 100644 --- a/src/mobius/integrations/_weight_loading.py +++ b/src/mobius/integrations/_weight_loading.py @@ -21,6 +21,7 @@ "apply_weights", "stream_qdq_safetensors_to_model", "stream_preprocessed_safetensors_to_model", + "stream_preprocessed_safetensors_to_package", "stream_safetensors_to_model", "external_data_checksums", ] @@ -77,6 +78,7 @@ class StreamingWeightSource: mode: Literal["direct", "fp8_scalar", "fp8_block_128"] = "direct" scale_name: str | None = None expected_scale: float | None = None + expected_dtype: str | None = None @dataclasses.dataclass(frozen=True) @@ -684,8 +686,8 @@ def tensor_func( ) -def stream_preprocessed_safetensors_to_model( - model: ir.Model, +def _stream_preprocessed_safetensors( + initializers: Mapping[str, ir.Value], model_id: str, planner: Callable[ [Mapping[str, tuple[str, list[int], str]], Mapping[str, ir.Value]], @@ -704,7 +706,7 @@ def stream_preprocessed_safetensors_to_model( """ paths = _resolve_shard_paths(model_id, revision) key_index = _shard_key_index(paths) - plan = planner(key_index, model.graph.initializers) + plan = planner(key_index, initializers) consumed = set(plan.ignored) | set(plan.constants) assigned: set[str] = set() @@ -725,6 +727,11 @@ def validate_source( if source.source_name not in key_index: raise ValueError(f"Streaming source '{source.source_name}' does not exist") _source_path, source_shape, source_dtype = key_index[source.source_name] + if source.expected_dtype is not None and source_dtype != source.expected_dtype: + raise ValueError( + f"Streaming source '{source.source_name}' has dtype {source_dtype}; " + f"expected {source.expected_dtype}" + ) if source.mode in {"fp8_block_128", "fp8_scalar"}: if source_dtype not in {"F8_E4M3", "F8_E5M2"}: raise ValueError( @@ -794,7 +801,7 @@ def validate_source( return source_shape, source_bytes, scale_bytes for target_name, source in plan.targets.items(): - initializer = model.graph.initializers.get(target_name) + initializer = initializers.get(target_name) if initializer is None: raise ValueError(f"Streaming plan targets unknown initializer '{target_name}'") if initializer.const_value is not None: @@ -882,7 +889,7 @@ def validate_source( missing_targets = sorted( name - for name, initializer in model.graph.initializers.items() + for name, initializer in initializers.items() if initializer.const_value is None and name not in assigned ) if missing_targets: @@ -916,10 +923,69 @@ def validate_source( "largest_reconstruction_working_set_bytes": (largest_reconstruction_working_set_bytes), **dict(plan.report), } + return report + + +def stream_preprocessed_safetensors_to_model( + model: ir.Model, + model_id: str, + planner: Callable[ + [Mapping[str, tuple[str, list[int], str]], Mapping[str, ir.Value]], + StreamingWeightPlan, + ], + *, + revision: str | None = None, +) -> dict[str, object]: + """Stream a fully classified transformed checkpoint into one dense ONNX graph.""" + report = _stream_preprocessed_safetensors( + model.graph.initializers, + model_id, + planner, + revision=revision, + ) model.metadata_props["mobius.weight_loading"] = json.dumps(report, sort_keys=True) return report +def stream_preprocessed_safetensors_to_package( + package: Mapping[str, ir.Model], + model_id: str, + planner: Callable[ + [Mapping[str, tuple[str, list[int], str]], Mapping[str, ir.Value]], + StreamingWeightPlan, + ], + *, + revision: str | None = None, +) -> dict[str, object]: + """Stream one fully classified checkpoint across a staged ONNX package. + + Initializer names are module-qualified and must be unique across every + component. Graph constants are excluded: the planner owns checkpoint + tensors, while fixed graph constants have no safetensors source. + """ + initializers: dict[str, ir.Value] = {} + for component_name, model in package.items(): + for name, initializer in model.graph.initializers.items(): + if initializer.const_value is not None: + continue + if name in initializers: + raise ValueError( + f"Staged streaming package has duplicate initializer {name!r} in " + f"component {component_name!r}." + ) + initializers[name] = initializer + report = _stream_preprocessed_safetensors( + initializers, + model_id, + planner, + revision=revision, + ) + serialized_report = json.dumps(report, sort_keys=True) + for model in package.values(): + model.metadata_props["mobius.weight_loading"] = serialized_report + return report + + def _graph_constant( graph: ir.Graph, name: str, diff --git a/src/mobius/integrations/_weight_loading_stream_test.py b/src/mobius/integrations/_weight_loading_stream_test.py index cf9396bdb..616ba8c4a 100644 --- a/src/mobius/integrations/_weight_loading_stream_test.py +++ b/src/mobius/integrations/_weight_loading_stream_test.py @@ -21,10 +21,14 @@ from onnx_ir import tensor_adapters from mobius._builder import build_from_module +from mobius._model_package import ModelPackage from mobius._testing import make_config from mobius.integrations._weight_loading import ( + StreamingWeightPlan, + StreamingWeightSource, _shard_key_index, external_data_checksums, + stream_preprocessed_safetensors_to_package, stream_safetensors_to_model, ) from mobius.models.base import CausalLMModel @@ -137,6 +141,33 @@ def test_assignment_is_deferred_not_materialized(self, tmp_path): ] assert lazy, "expected streamed weights to be deferred LazyTensors" + def test_preprocessed_package_streams_and_records_every_component( + self, tmp_path, monkeypatch + ): + def _no_hub(*_a, **_k): + raise AssertionError("streaming a local dir must not call the Hub") + + monkeypatch.setattr("mobius.integrations._weight_loading.hf_hub_download", _no_hub) + model = _fresh_model() + state = _make_checkpoint_state(model) + _save_single(state, tmp_path) + package = ModelPackage({"decoder": model}) + + def plan(sources, initializers): + return StreamingWeightPlan( + targets={ + name: StreamingWeightSource(name, expected_dtype="F32") + for name, initializer in initializers.items() + if initializer.const_value is None + } + ) + + report = stream_preprocessed_safetensors_to_package(package, str(tmp_path), plan) + + assert report["assigned_tensors"] == len(state) + assert model.metadata_props["mobius.weight_loading"] + _roundtrip_and_compare(model, state, tmp_path) + class TestStreamingRefusals: def test_duplicate_tensor_across_shards_is_rejected(self, tmp_path): diff --git a/src/mobius/integrations/gguf/_builder.py b/src/mobius/integrations/gguf/_builder.py index d7d6a941a..c58fc5d79 100644 --- a/src/mobius/integrations/gguf/_builder.py +++ b/src/mobius/integrations/gguf/_builder.py @@ -480,6 +480,7 @@ def _gguf_header_info_from_header_prefix( data, source=source, require_architecture=False, + collect_tensor_type_ids=True, ) @@ -598,6 +599,15 @@ def read_response(response) -> list[bytes]: ) return _GGUFPreflightFallbackRevision(commit_hash) _validate_preflight_split_header(header_info, source=source) + from mobius.integrations._vibeasr_bitnet import reject_vibeasr_bitnet_gguf + + reject_vibeasr_bitnet_gguf( + source=source, + repository=repo_id, + revision=commit_hash, + filename=filename, + header=header_info, + ) architecture = header_info.architecture if ( dispatch_architecture @@ -665,6 +675,29 @@ def _validate_gguf_model( if not isinstance(gguf_model, GgufShardSet): split_count = int(gguf_model.get_metadata("split.count", 1)) _raise_for_sharded_gguf(source=source, split_count=split_count) + from mobius.integrations._vibeasr_bitnet import reject_vibeasr_bitnet_gguf + + reader_tensors = getattr(gguf_model, "reader_tensors", None) + tensor_type_ids = frozenset() + if callable(reader_tensors): + tensor_type_ids = frozenset( + int(getattr(tensor.tensor_type, "value", tensor.tensor_type)) + for tensor in reader_tensors() + ) + reject_vibeasr_bitnet_gguf( + source=source, + header=GGUFHeaderInfo( + architecture=gguf_model.architecture, + tensor_count=len(gguf_model.tensor_names), + split_no=None, + split_count=None, + split_tensors_count=None, + name=gguf_model.get_metadata("general.name", None), + file_type=gguf_model.get_metadata("general.file_type", None), + quantization_version=gguf_model.get_metadata("general.quantization_version", None), + tensor_type_ids=tensor_type_ids, + ), + ) from mobius.integrations.gguf._qwen4_exp import validate_qwen4exp_tensor_contract validate_qwen4exp_tensor_contract( diff --git a/src/mobius/integrations/gguf/_errors.py b/src/mobius/integrations/gguf/_errors.py index 6c761197b..7ddc43014 100644 --- a/src/mobius/integrations/gguf/_errors.py +++ b/src/mobius/integrations/gguf/_errors.py @@ -24,6 +24,7 @@ "ShardedGGUFNotSupportedError", "UnsupportedGGUFArchitectureError", "UnsupportedGGUFQuantizationError", + "VibeASRBitNetGGUFImportError", ] @@ -39,6 +40,10 @@ class UnsupportedGGUFQuantizationError(ValueError): """A stored GGML tensor type that mobius cannot read or preserve.""" +class VibeASRBitNetGGUFImportError(UnsupportedGGUFQuantizationError): + """A VibeASR.cpp-native GGUF whose execution contract has no ORT equivalent.""" + + class DisabledGGUFArchitectureError(NotImplementedError): """A GGUF architecture whose conversion is deliberately turned off. diff --git a/src/mobius/integrations/gguf/_header.py b/src/mobius/integrations/gguf/_header.py index 2fcf6f456..f98def1e1 100644 --- a/src/mobius/integrations/gguf/_header.py +++ b/src/mobius/integrations/gguf/_header.py @@ -10,6 +10,13 @@ from typing import Any _GGUF_ARCHITECTURE_KEY = b"general.architecture" +_GGUF_IDENTITY_STRING_KEYS = { + b"general.name": "name", +} +_GGUF_IDENTITY_INTEGER_KEYS = { + b"general.file_type": "file_type", + b"general.quantization_version": "quantization_version", +} _GGUF_SPLIT_KEYS = { b"split.no": "split_no", b"split.count": "split_count", @@ -58,6 +65,10 @@ class GGUFHeaderInfo: split_no: int | None split_count: int | None split_tensors_count: int | None + name: str | None = None + file_type: int | None = None + quantization_version: int | None = None + tensor_type_ids: frozenset[int] = frozenset() def _gguf_header_info_from_header( @@ -65,6 +76,7 @@ def _gguf_header_info_from_header( *, source: str, require_architecture: bool = True, + collect_tensor_type_ids: bool = False, ) -> GGUFHeaderInfo: """Validate a GGUF metadata table and return bounded preflight fields.""" size = len(data) @@ -95,6 +107,15 @@ def read_uint64(offset: int) -> tuple[int, int]: raise GGUFHeaderTruncatedError(f"{source!r} has a truncated GGUF metadata header.") return struct.unpack_from(f"{byte_order}Q", data, offset)[0], end + def skip_bytes(offset: int, count: int, *, field_name: str) -> int: + end = offset + count + if end > size: + raise GGUFHeaderTruncatedError( + f"{source!r} has a truncated GGUF {field_name}: " + f"requires {count} bytes with only {size - offset} remaining." + ) + return end + def read_string_span(offset: int) -> tuple[int, int, int]: length, offset = read_uint64(offset) end = offset + length @@ -105,18 +126,17 @@ def read_string_span(offset: int) -> tuple[int, int, int]: ) return offset, end, end - def read_integer(value_type: int, offset: int) -> tuple[int, int]: + def read_integer(value_type: int, offset: int, *, key_name: str) -> tuple[int, int]: format_char = _GGUF_INTEGER_FORMATS.get(value_type) if format_char is None: raise ValueError( - f"{source!r} encodes split bookkeeping with non-integer GGUF " - f"type {value_type}." + f"{source!r} encodes {key_name} with non-integer GGUF type {value_type}." ) width = _GGUF_SCALAR_WIDTHS[value_type] end = offset + width if end > size: raise GGUFHeaderTruncatedError( - f"{source!r} has a truncated GGUF split metadata value." + f"{source!r} has a truncated GGUF {key_name} metadata value." ) return int(struct.unpack_from(f"{byte_order}{format_char}", data, offset)[0]), end @@ -188,6 +208,12 @@ def skip_value(value_type: int, offset: int, *, depth: int = 0) -> int: split_values: dict[str, list[int]] = { field_name: [] for field_name in _GGUF_SPLIT_KEYS.values() } + identity_string_values: dict[str, list[bytes]] = { + field_name: [] for field_name in _GGUF_IDENTITY_STRING_KEYS.values() + } + identity_integer_values: dict[str, list[int]] = { + field_name: [] for field_name in _GGUF_IDENTITY_INTEGER_KEYS.values() + } for _ in range(kv_count): key_start, key_end, offset = read_string_span(offset) value_type, offset = read_uint32(offset) @@ -204,8 +230,25 @@ def skip_value(value_type: int, offset: int, *, depth: int = 0) -> int: ) value_start, value_end, offset = read_string_span(offset) architecture_values.append(bytes(data[value_start:value_end])) + elif (field_name := _GGUF_IDENTITY_STRING_KEYS.get(key)) is not None: + if value_type != _GGUF_STRING: + raise ValueError( + f"{source!r} encodes general.{field_name} with GGUF type " + f"{value_type}, expected string type {_GGUF_STRING}." + ) + value_start, value_end, offset = read_string_span(offset) + identity_string_values[field_name].append(bytes(data[value_start:value_end])) + elif (field_name := _GGUF_IDENTITY_INTEGER_KEYS.get(key)) is not None: + value, offset = read_integer( + value_type, + offset, + key_name=f"general.{field_name}", + ) + identity_integer_values[field_name].append(value) elif (field_name := _GGUF_SPLIT_KEYS.get(key)) is not None: - value, offset = read_integer(value_type, offset) + value, offset = read_integer( + value_type, offset, key_name=field_name.replace("_", ".") + ) split_values[field_name].append(value) else: offset = skip_value(value_type, offset) @@ -232,12 +275,49 @@ def skip_value(value_type: int, offset: int, *, depth: int = 0) -> int: f"{source!r} contains duplicate {field_name.replace('_', '.')} metadata." ) split_fields[field_name] = values[0] if values else None + identity_strings: dict[str, str | None] = {} + for field_name, values in identity_string_values.items(): + if len(values) > 1: + raise ValueError(f"{source!r} contains duplicate general.{field_name} metadata.") + if not values: + identity_strings[field_name] = None + continue + try: + identity_strings[field_name] = values[0].decode("utf-8") + except UnicodeDecodeError as error: + raise ValueError( + f"{source!r} has a non-UTF-8 general.{field_name} value." + ) from error + identity_integers: dict[str, int | None] = {} + for field_name, values in identity_integer_values.items(): + if len(values) > 1: + raise ValueError(f"{source!r} contains duplicate general.{field_name} metadata.") + identity_integers[field_name] = values[0] if values else None + tensor_type_ids: frozenset[int] = frozenset() + if collect_tensor_type_ids: + types: set[int] = set() + for tensor_index in range(tensor_count): + _, _, offset = read_string_span(offset) + dimensions, offset = read_uint32(offset) + offset = skip_bytes( + offset, + dimensions * 8, + field_name=f"tensor {tensor_index} dimensions", + ) + tensor_type, offset = read_uint32(offset) + types.add(tensor_type) + offset = skip_bytes(offset, 8, field_name=f"tensor {tensor_index} offset") + tensor_type_ids = frozenset(types) return GGUFHeaderInfo( architecture=architecture, tensor_count=tensor_count, split_no=split_fields["split_no"], split_count=split_fields["split_count"], split_tensors_count=split_fields["split_tensors_count"], + name=identity_strings["name"], + file_type=identity_integers["file_type"], + quantization_version=identity_integers["quantization_version"], + tensor_type_ids=tensor_type_ids, ) diff --git a/src/mobius/integrations/gguf/_preflight.py b/src/mobius/integrations/gguf/_preflight.py index 89a9999d6..21e58f5b4 100644 --- a/src/mobius/integrations/gguf/_preflight.py +++ b/src/mobius/integrations/gguf/_preflight.py @@ -605,8 +605,8 @@ def preflight_hf_gguf( ) ) - architecture, num_experts, total_tensors, total_params = _hf_gguf_metadata( - api, repo_id, revision, token + architecture, num_experts, total_tensors, total_params, resolved_revision = ( + _hf_gguf_metadata(api, repo_id, revision, token) ) model_type = resolve_model_type(architecture) if architecture else None quantization = _detect_quantization(filename, *shard_files) @@ -618,6 +618,22 @@ def preflight_hf_gguf( quantization=quantization, source=f"{repo_id}{('/' + filename) if filename else ''}", ) + from mobius.integrations._vibeasr_bitnet import find_vibeasr_bitnet_gguf_artifact + + for file in files: + artifact = find_vibeasr_bitnet_gguf_artifact( + repository=repo_id, + revision=resolved_revision or revision, + filename=file.filename, + size_bytes=file.size_bytes, + sha256=file.sha256, + ) + if artifact is not None: + blockers.append( + f"VibeVoice ASR BitNet native GGUF blocker for {artifact.filename}: " + f"{artifact.blocker} Build from the dense F32 safetensors conversion source " + "instead; Mobius does not claim native BitNet/GGUF execution." + ) warnings: list[str] = [] if architecture is None: @@ -754,10 +770,10 @@ def _shards_for_prefix(gguf_files: list[str], prefix: str) -> list[str]: def _hf_gguf_metadata( api: Any, repo_id: str, revision: str | None, token: str | bool | None -) -> tuple[str | None, int | None, int | None, int | None]: - """Return ``(architecture, num_experts, total_tensors, total_params)``. +) -> tuple[str | None, int | None, int | None, int | None, str | None]: + """Return architecture/count metadata plus the immutable resolved revision. - Uses ``model_info(expand=["gguf"])`` which surfaces the parsed GGUF header + Uses ``model_info(expand=["gguf", "sha"])`` which surfaces the parsed GGUF header fields without downloading tensor bytes. The Hub's ``gguf.total`` is the *parameter* count (not the tensor count); any field the Hub does not expose comes back ``None`` (the caller records a warning). @@ -766,12 +782,15 @@ def _hf_gguf_metadata( num_experts: int | None = None total_tensors: int | None = None total_params: int | None = None + resolved_revision: str | None = None try: - info = api.model_info(repo_id, revision=revision, token=token, expand=["gguf"]) + info = api.model_info(repo_id, revision=revision, token=token, expand=["gguf", "sha"]) except Exception as error: logger.info("model_info(expand=gguf) unavailable for %s: %s", repo_id, error) - return architecture, num_experts, total_tensors, total_params + return architecture, num_experts, total_tensors, total_params, resolved_revision + sha = getattr(info, "sha", None) + resolved_revision = sha if isinstance(sha, str) else None gguf_meta = getattr(info, "gguf", None) if isinstance(gguf_meta, dict): architecture = gguf_meta.get("architecture") or gguf_meta.get("general.architecture") @@ -792,6 +811,7 @@ def _hf_gguf_metadata( int(num_experts) if num_experts else None, int(total_tensors) if total_tensors else None, int(total_params) if total_params else None, + resolved_revision, ) diff --git a/src/mobius/integrations/gguf/_reader.py b/src/mobius/integrations/gguf/_reader.py index 6866162e6..baa9e42cf 100644 --- a/src/mobius/integrations/gguf/_reader.py +++ b/src/mobius/integrations/gguf/_reader.py @@ -38,7 +38,7 @@ import numpy as np -from mobius.integrations.gguf._header import _gguf_architecture_from_header +from mobius.integrations.gguf._header import _gguf_header_info_from_header logger = logging.getLogger(__name__) @@ -247,11 +247,15 @@ def __init__(self, path: str | Path, *, follow_symlinks: bool = True) -> None: ) with os.fdopen(os.dup(descriptor), "rb") as stream: with mmap.mmap(stream.fileno(), length=0, access=mmap.ACCESS_READ) as mapped: - _gguf_architecture_from_header( + header_info = _gguf_header_info_from_header( mapped, source=str(self._path), require_architecture=False, + collect_tensor_type_ids=True, ) + from mobius.integrations._vibeasr_bitnet import reject_vibeasr_bitnet_gguf + + reject_vibeasr_bitnet_gguf(source=str(self._path), header=header_info) stream.seek(0) self._reader = GGUFReader(cast(Any, stream)) stream.seek(0) diff --git a/src/mobius/integrations/transformers/_builder.py b/src/mobius/integrations/transformers/_builder.py index 4672b6737..c0c7733d0 100644 --- a/src/mobius/integrations/transformers/_builder.py +++ b/src/mobius/integrations/transformers/_builder.py @@ -22,6 +22,7 @@ from mobius.integrations._weight_loading import ( _download_weights, stream_preprocessed_safetensors_to_model, + stream_preprocessed_safetensors_to_package, stream_qdq_safetensors_to_model, ) from mobius.integrations.compressed_tensors import ( @@ -355,6 +356,26 @@ def build_transformers_model( # processor contract. Keep config detection and weight loading pinned. revision = VIBEVOICE_ASR_REVISION detection_revision = VIBEVOICE_ASR_REVISION + + from mobius.integrations._vibeasr_bitnet import ( + VIBEVOICE_ASR_BITNET_REPOSITORY, + VIBEVOICE_ASR_BITNET_REVISION, + build_vibeasr_bitnet_dense_weight_plan, + is_vibeasr_bitnet_conversion_source, + ) + + if model_id == VIBEVOICE_ASR_BITNET_REPOSITORY: + if detection_revision is None: + # The release mixes dense F32 conversion weights with VibeASR.cpp-native + # GGUFs. Pin config and weight retrieval to the audited dense source. + revision = VIBEVOICE_ASR_BITNET_REVISION + detection_revision = VIBEVOICE_ASR_BITNET_REVISION + elif detection_revision != VIBEVOICE_ASR_BITNET_REVISION: + raise ValueError( + f"{VIBEVOICE_ASR_BITNET_REPOSITORY} is supported only at the audited " + f"revision {VIBEVOICE_ASR_BITNET_REVISION}; got {detection_revision!r}." + ) + is_vibeasr_bitnet_dense_source = is_vibeasr_bitnet_conversion_source(model_id, revision) if model_id == "nvidia/RE-USE" and detection_revision is None: # Pin the very first AutoConfig/raw-JSON probe, not only the later # bespoke loader. Otherwise mutable Hub main could change dispatch @@ -571,13 +592,35 @@ def build_transformers_model( ) for name, model in package.items(): model.graph.name = f"{model_id}/{name}" - if model_type in _QWEN4_MODEL_TYPES | {"vibevoice", "VibeVoiceForASRTraining"}: + if ( + model_type + in _QWEN4_MODEL_TYPES + | { + "vibevoice", + "VibeVoiceForASRTraining", + } + or is_vibeasr_bitnet_dense_source + ): model.metadata_props["mobius.source_revision"] = revision or "unpinned" if load_weights: - _reject_unsupported_affine_qwen4(model_type, config) - if config.block_quant_scheme is not None and hasattr( - model_module, "build_fp8_streaming_plan" + if is_vibeasr_bitnet_dense_source: + package.weight_loading_report = stream_preprocessed_safetensors_to_package( + package, + model_id, + lambda source_tensors, initializers: build_vibeasr_bitnet_dense_weight_plan( + model_module, + source_tensors, + initializers, + ), + revision=revision, + ) + else: + _reject_unsupported_affine_qwen4(model_type, config) + if ( + not is_vibeasr_bitnet_dense_source + and config.block_quant_scheme is not None + and hasattr(model_module, "build_fp8_streaming_plan") ): if len(package) != 1: raise ValueError( @@ -608,7 +651,7 @@ def build_transformers_model( "native FP8 was not preserved. See weight-loading-report.json.", model_id, ) - elif model_type in _QWEN4_MODEL_TYPES: + elif not is_vibeasr_bitnet_dense_source and model_type in _QWEN4_MODEL_TYPES: from mobius.integrations.transformers._qwen4_exp_weights import ( stream_qwen4_exp_safetensors_to_package, ) @@ -619,7 +662,7 @@ def build_transformers_model( config, revision=revision, ) - elif compressed_tensors_config is not None: + elif not is_vibeasr_bitnet_dense_source and compressed_tensors_config is not None: stream_compressed_tensors_to_package( package, model_id, @@ -629,7 +672,7 @@ def build_transformers_model( fp8_kv_cache=fp8_kv_cache, keep_quantized=keep_quantized, ) - else: + elif not is_vibeasr_bitnet_dense_source: state_dict = _download_weights(model_id, revision=revision) if hasattr(model_module, "preprocess_weights"): state_dict = model_module.preprocess_weights(state_dict)