From 335ed97d12027270743388438ed8c7aaeb16058d Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Wed, 26 Aug 2026 19:31:15 -0700 Subject: [PATCH 1/3] Load quantized weights per component Parse authoritative component quantization plans, map component-local modules to HuggingFace names, honor exact and regex float exclusions per projection, and normalize existing packed sidecars through typed codecs. Mobius does not quantize float weights. Signed-off-by: Xiaoyu Zhang --- CHANGELOG.md | 13 + src/mobius/__main__.py | 29 +- src/mobius/_builder.py | 12 +- src/mobius/_component_quantization.py | 574 ++++++++++++++++++ src/mobius/_component_quantization_test.py | 268 ++++++++ src/mobius/_configs/__init__.py | 3 +- src/mobius/_configs/_base.py | 134 ++++ src/mobius/_configs/_quantization.py | 147 ++++- src/mobius/_configs_test.py | 71 +++ src/mobius/components/__init__.py | 4 + src/mobius/components/_quantized_linear.py | 75 +++ .../components/_quantized_linear_test.py | 46 ++ src/mobius/components/_vision.py | 12 +- .../integrations/transformers/_builder.py | 19 + .../transformers/_builder_test.py | 25 + src/mobius/models/gemma4_test.py | 33 + src/mobius/models/mage_vl.py | 3 +- tests/build_graph_test.py | 85 +++ 18 files changed, 1522 insertions(+), 31 deletions(-) create mode 100644 src/mobius/_component_quantization.py create mode 100644 src/mobius/_component_quantization_test.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 079bc5a62..3a56b6a2e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,19 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Per-component quantized checkpoint loading + +#### Added + +- Multi-component checkpoints may declare an authoritative + `component_quantization` mapping with independent affine layouts for decoder, + encoder, vision, audio, and embedding components. Exact and regex + `modules_to_not_convert` rules are evaluated for each component-local module + against its HuggingFace source name, so selected projections remain floating + point while the rest of the component binds existing packed weights. +- Mobius validates and normalizes existing Olive, GPTQ, and AWQ sidecars per + component. It does not quantize floating-point checkpoint weights. + ### Packed fused MoE experts (Olive/GPTQ/AWQ) survive HF weight renaming #### Fixed diff --git a/src/mobius/__main__.py b/src/mobius/__main__.py index 05c477064..6d4011acc 100644 --- a/src/mobius/__main__.py +++ b/src/mobius/__main__.py @@ -389,7 +389,12 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask: ) compressed_tensors_config = CompressedTensorsConfig.from_hf_config(parent_config) - config = _config_from_hf(hf_config, parent_config=parent_config) + module_class = registry.get(model_type) + config = _config_from_hf( + hf_config, + parent_config=parent_config, + module_class=module_class, + ) if dtype_override is not None: config = dataclasses.replace(config, dtype=dtype_override) elif compressed_tensors_config is not None and keep_quantized: @@ -425,7 +430,14 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask: task = _resolve_static_cache_task(model_type) elif task is None: task = _default_task_for_model(model_type) - module_class = registry.get(model_type) + from mobius.tasks import get_task + + resolved_task = get_task(task) + component_manifest = resolved_task.component_manifest( + module_class=module_class, + model_type=model_type, + hf_config=parent_config, + ) model_module = module_class(config) pkg = build_from_module( model_module, @@ -435,6 +447,7 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask: fp8_kv_cache=fp8_kv_cache, kv_cache_scales=kv_cache_scales, prune_prefill_prefix=prune_prefill_prefix, + component_manifest=component_manifest, ) for name, model in pkg.items(): model.graph.name = f"{config_path}/{name}" @@ -464,6 +477,18 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask: state_dict = _load_weights_from_dir(config_path) if hasattr(model_module, "preprocess_weights"): state_dict = model_module.preprocess_weights(state_dict) + from mobius._component_quantization import ( + normalize_component_quantized_weights, + ) + + state_dict = normalize_component_quantized_weights( + state_dict, + model_module, + config, + pkg.keys(), + manifest=component_manifest, + task=resolved_task, + ) pkg.apply_weights(state_dict) else: model_id_or_path = args.model diff --git a/src/mobius/_builder.py b/src/mobius/_builder.py index 69f1a6589..5292cca26 100644 --- a/src/mobius/_builder.py +++ b/src/mobius/_builder.py @@ -25,6 +25,8 @@ from onnxscript import nn from mobius._build_context import build_context +from mobius._component_manifest import ComponentManifest +from mobius._component_quantization import configure_component_quantization from mobius._configs import BaseModelConfig from mobius._execution_providers import ep_registry from mobius._flags import flags @@ -136,6 +138,7 @@ def build_from_module( fp8_kv_cache: bool = False, kv_cache_scales: dict[int, tuple[float, float]] | None = None, prune_prefill_prefix: bool = False, + component_manifest: ComponentManifest | None = None, ) -> ModelPackage: """Build an ONNX :class:`ModelPackage` from a module instance and config. @@ -158,11 +161,16 @@ def build_from_module( if hasattr(config, "validate"): config.validate() dtype = getattr(config, "dtype", ir.DataType.FLOAT) - _cast_module_dtype(module, dtype) if prune_prefill_prefix: task = _enable_prefill_prefix_pruning_task(task) resolved_task = get_task(task) - component_manifest = resolved_task.component_manifest() + component_manifest = configure_component_quantization( + module, + config, + resolved_task, + manifest=component_manifest, + ) + _cast_module_dtype(module, dtype) capabilities = ep_registry.require(execution_provider) with build_context(capabilities, dtype): package = resolved_task.build(module, config) diff --git a/src/mobius/_component_quantization.py b/src/mobius/_component_quantization.py new file mode 100644 index 000000000..eee210cfe --- /dev/null +++ b/src/mobius/_component_quantization.py @@ -0,0 +1,574 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Configure and load independently quantized model-package components.""" + +from __future__ import annotations + +__all__ = [ + "configure_component_quantization", + "normalize_component_quantized_weights", +] + +from collections.abc import Iterable, Mapping +from typing import Any + +import onnx_ir as ir +from onnxscript import nn + +from mobius._component_manifest import ComponentDescriptor, ComponentManifest +from mobius._configs import BaseModelConfig, QuantizationConfig +from mobius.components import ( + ClippableLinear, + ClippableQuantizedLinear, + Embedding, + Linear, + QuantizedEmbedding, + QuantizedLinear, + make_clippable_quantized_linear_factory, + make_quantized_linear_factory, +) +from mobius.tasks import ModelTask, get_task +from mobius.weights import FloatWeight, PackedWeight, codec_registry + +_AFFINE_METHODS = frozenset({"olive", "gptq", "awq"}) +_KNOWN_SPLIT_COMPONENTS = frozenset( + { + "decoder", + "encoder", + "vision", + "vision_encoder", + "audio", + "audio_encoder", + "embedding", + "model", + } +) + + +def _resolve_module(root: nn.Module, path: str) -> nn.Module | None: + if not path: + return root + current: object = root + for part in path.split("."): + if not hasattr(current, part): + return None + current = getattr(current, part) + return current if isinstance(current, nn.Module) else None + + +def _replace_child(root: nn.Module, path: str, replacement: nn.Module) -> None: + if not path: + raise ValueError("Cannot replace a component's root module") + parts = path.split(".") + parent: object = root + for part in parts[:-1]: + parent = getattr(parent, part) + old = getattr(parent, parts[-1]) + if hasattr(replacement, "_set_name") and hasattr(old, "name"): + replacement._set_name(old.name) + setattr(parent, parts[-1], replacement) + + +def _component_quantization( + config: BaseModelConfig, + component: str, +) -> QuantizationConfig | None: + resolver = getattr(config, "quantization_for", None) + if resolver is not None: + return resolver(component) + mapping = getattr(config, "component_quantization", None) + if mapping is None: + return getattr(config, "quantization", None) + if component in mapping: + return mapping[component] + if component == "model": + return mapping.get("decoder") + if component == "decoder": + return mapping.get("model") + return None + + +def _linear_factory( + config: BaseModelConfig, + quantization: QuantizationConfig, +) -> type[nn.Module]: + zero_point_dtype = config.dtype if quantization.float_zero_point else ir.DataType.UINT8 + return make_quantized_linear_factory( + bits=quantization.bits, + block_size=quantization.group_size, + has_zero_point=not quantization.sym, + zero_point_dtype=zero_point_dtype, + ) + + +def _clippable_linear_factory( + config: BaseModelConfig, + quantization: QuantizationConfig, +) -> type[nn.Module]: + zero_point_dtype = config.dtype if quantization.float_zero_point else ir.DataType.UINT8 + return make_clippable_quantized_linear_factory( + bits=quantization.bits, + block_size=quantization.group_size, + has_zero_point=not quantization.sym, + zero_point_dtype=zero_point_dtype, + ) + + +def _float_linear(module: QuantizedLinear) -> Linear: + return Linear(module._k, module._n, bias=module.bias is not None) + + +def _float_embedding(module: QuantizedEmbedding) -> Embedding: + return Embedding( + int(module.qweight.shape[0]), + module._embedding_dim, + module.padding_idx, + ) + + +def _effective_module_quantization( + component_quantization: QuantizationConfig | None, + descriptor: ComponentDescriptor, + local_module_path: str, +) -> QuantizationConfig | None: + if component_quantization is None or component_quantization.quant_method == "none": + return None + source_names = descriptor.source_module_names(local_module_path) + return component_quantization.for_module(source_names) + + +def _configure_component_module( + component_module: nn.Module, + descriptor: ComponentDescriptor, + config: BaseModelConfig, + component_quantization: QuantizationConfig | None, + *, + owned_by_other_components: tuple[str, ...] = (), +) -> None: + replacements: list[tuple[str, nn.Module]] = [] + for local_path, child in list(component_module.named_modules()): + if not local_path: + continue + if any( + local_path == prefix or local_path.startswith(f"{prefix}.") + for prefix in owned_by_other_components + ): + continue + quantization = _effective_module_quantization( + component_quantization, + descriptor, + local_path, + ) + is_lm_head = local_path == "lm_head" or local_path.endswith(".lm_head") + if quantization is not None and is_lm_head and not quantization.quantize_lm_head: + quantization = None + + if isinstance(child, ClippableQuantizedLinear): + if type(child).forward is not ClippableQuantizedLinear.forward: + raise TypeError( + f"Component plan cannot rewrite specialized clipped " + f"quantized module {local_path!r} " + f"({type(child).__name__}); provide a model weight adapter " + "for this component." + ) + replacement: nn.Module + if quantization is None: + replacement = ClippableLinear( + child._k, + child._n, + bias=child.bias is not None, + ) + else: + replacement = _clippable_linear_factory(config, quantization)( + child._k, + child._n, + bias=child.bias is not None, + ) + replacements.append((local_path, replacement)) + continue + + if isinstance(child, QuantizedLinear): + if type(child).forward is not QuantizedLinear.forward: + raise TypeError( + f"Component plan cannot rewrite specialized quantized " + f"module {local_path!r} ({type(child).__name__}); provide " + "a model weight adapter for this component." + ) + replacement = ( + _float_linear(child) + if quantization is None + else _linear_factory(config, quantization)( + child._k, + child._n, + bias=child.bias is not None, + ) + ) + replacements.append((local_path, replacement)) + continue + + if isinstance(child, QuantizedEmbedding): + if type(child).forward is not QuantizedEmbedding.forward: + raise TypeError( + f"Component plan cannot rewrite specialized quantized " + f"embedding {local_path!r} ({type(child).__name__}); " + "provide a model weight adapter for this component." + ) + if quantization is None or not quantization.quantize_embeddings: + replacements.append((local_path, _float_embedding(child))) + continue + + if quantization is None: + continue + if quantization.quant_method not in _AFFINE_METHODS: + continue + + if isinstance(child, Embedding) and type(child).forward is Embedding.forward: + embedding_dim = int(child.weight.shape[1]) + if ( + quantization.quantize_embeddings + and embedding_dim % quantization.group_size == 0 + ): + num_embeddings = int(child.weight.shape[0]) + replacements.append( + ( + local_path, + QuantizedEmbedding( + num_embeddings, + embedding_dim, + bits=quantization.bits, + block_size=quantization.group_size, + has_zero_point=not quantization.sym, + padding_idx=child.padding_idx, + ), + ) + ) + continue + + if isinstance(child, Linear) and type(child).forward is Linear.forward: + out_features, in_features = (int(dim) for dim in child.weight.shape) + replacements.append( + ( + local_path, + _linear_factory(config, quantization)( + in_features, + out_features, + bias=child.bias is not None, + ), + ) + ) + elif type(child) is ClippableLinear: + out_features, in_features = (int(dim) for dim in child.weight.shape) + replacements.append( + ( + local_path, + _clippable_linear_factory(config, quantization)( + in_features, + out_features, + bias=child.bias is not None, + ), + ) + ) + + # Replace deepest children first so replacing a parent never invalidates a + # path that still needs to be visited. + for path, replacement in sorted( + replacements, + key=lambda item: item[0].count("."), + reverse=True, + ): + _replace_child(component_module, path, replacement) + + +def _default_manifest( + module: nn.Module, + config: BaseModelConfig, + task: str | ModelTask, +) -> ComponentManifest: + resolved_task = get_task(task) + model_type = getattr(config, "model_type", None) + return resolved_task.component_manifest( + module_class=type(module), + model_type=model_type, + hf_config=config, + ) + + +def configure_component_quantization( + module: nn.Module, + config: BaseModelConfig, + task: str | ModelTask, + *, + manifest: ComponentManifest | None = None, +) -> ComponentManifest: + """Apply authoritative component plans to graph parameter scaffolding.""" + manifest = manifest or _default_manifest(module, config, task) + mapping = getattr(config, "component_quantization", None) + root_quantization = getattr(config, "quantization", None) + single_component_rules = ( + manifest.names == ("model",) + and root_quantization is not None + and (root_quantization.modules_to_not_convert or root_quantization.overrides) + ) + if mapping is None and not single_component_rules: + return manifest + + unresolved = set(mapping) - set(manifest) + if "model" in manifest: + unresolved.discard("decoder") + if "decoder" in manifest: + unresolved.discard("model") + if manifest.names == ("model",): + unresolved -= _KNOWN_SPLIT_COMPONENTS + if unresolved: + raise ValueError( + f"Component quantization references components not produced by " + f"{type(get_task(task)).__name__}: {sorted(unresolved)}. " + f"Available components: {sorted(manifest)}" + ) + + for descriptor in manifest.values(): + component_module = _resolve_module( + module, + descriptor.module_attribute_path, + ) + quantization = _component_quantization(config, descriptor.name) + if component_module is None: + continue + owned_elsewhere = ( + tuple( + other.module_attribute_path + for other in manifest.values() + if other.name != descriptor.name + and other.module_attribute_path + ) + if not descriptor.module_attribute_path + else () + ) + _configure_component_module( + component_module, + descriptor, + config, + quantization, + owned_by_other_components=owned_elsewhere, + ) + return manifest + + +def _raw_qweight_key(name: str) -> bool: + return name.endswith(("_qweight", ".qweight")) + + +def _canonical_component_parameter_keys( + module: nn.Module, + descriptor: ComponentDescriptor, +) -> frozenset[str]: + component_module = _resolve_module( + module, + descriptor.module_attribute_path, + ) + if component_module is None: + return frozenset() + + keys: set[str] = set() + prefixes = { + prefix + for prefix in (descriptor.name, descriptor.module_attribute_path) + if prefix + } + for local_path, child in component_module.named_modules(): + if not local_path: + continue + stems = { + f"{prefix}.{local_path}" if prefix else local_path for prefix in (*prefixes, "") + } + if isinstance(child, QuantizedEmbedding): + for stem in stems: + keys.update( + { + f"{stem}.qweight", + f"{stem}.scales", + f"{stem}.zero_points", + } + ) + elif isinstance(child, QuantizedLinear): + for stem in stems: + keys.update( + { + f"{stem}.weight", + f"{stem}.scales", + f"{stem}.zero_points", + } + ) + return frozenset(keys) + + +def _route_component_weights( + state_dict: Mapping[str, Any], + manifest: ComponentManifest, + component_names: tuple[str, ...], +) -> dict[str, dict[str, Any]]: + if len(component_names) == 1: + return {component_names[0]: dict(state_dict)} + + prefixes = { + name: { + prefix + for prefix in ( + name, + manifest[name].module_attribute_path, + ) + if prefix + } + for name in component_names + } + + def owner(key: str) -> str | None: + matches = [ + (len(prefix), component) + for component, component_prefixes in prefixes.items() + for prefix in component_prefixes + if key.startswith(f"{prefix}.") + ] + if not matches: + root_components = [ + name + for name in component_names + if not manifest[name].module_attribute_path + ] + return root_components[0] if len(root_components) == 1 else None + max_length = max(length for length, _ in matches) + owners = {component for length, component in matches if length == max_length} + if len(owners) != 1: + raise ValueError( + f"Checkpoint weight {key!r} matches multiple components " + f"{sorted(owners)} at the same prefix depth" + ) + return next(iter(owners)) + + routed = {name: {} for name in component_names} + for key, value in state_dict.items(): + component = owner(key) + if component is not None: + routed[component][key] = value + return routed + + +def _local_weight_module_path( + record_name: str, + descriptor: ComponentDescriptor, +) -> str: + name = record_name.removesuffix(".weight") + for prefix in (descriptor.module_attribute_path, descriptor.name): + if prefix and name.startswith(f"{prefix}."): + return name[len(prefix) + 1 :] + return name + + +def normalize_component_quantized_weights( + state_dict: dict[str, Any], + module: nn.Module, + config: BaseModelConfig, + component_names: Iterable[str], + *, + manifest: ComponentManifest | None = None, + task: str | ModelTask, +) -> dict[str, Any]: + """Normalize existing packed sidecars with each component's own plan.""" + component_names = tuple(component_names) + mapping = getattr(config, "component_quantization", None) + root_quantization = getattr(config, "quantization", None) + single_component_rules = ( + len(component_names) == 1 + and root_quantization is not None + and (root_quantization.modules_to_not_convert or root_quantization.overrides) + ) + if mapping is None and not single_component_rules: + return state_dict + manifest = manifest or _default_manifest(module, config, task) + manifest = manifest or _default_manifest(module, config, task) + routed = _route_component_weights(state_dict, manifest, component_names) + result = dict(state_dict) + + for component in component_names: + weights = routed[component] + canonical_keys = _canonical_component_parameter_keys( + module, + manifest[component], + ) + source_weights = { + key: value for key, value in weights.items() if key not in canonical_keys + } + if not any(_raw_qweight_key(key) for key in source_weights): + continue + + descriptor = manifest[component] + component_quantization = _component_quantization(config, component) + if component_quantization is None: + packed_key = next(key for key in weights if _raw_qweight_key(key)) + raise ValueError( + f"Component {component!r} is floating point but checkpoint " + f"contains packed weight {packed_key!r}" + ) + if component_quantization.quant_method not in codec_registry: + raise KeyError( + f"No packed-weight codec for component {component!r} method " + f"{component_quantization.quant_method!r}" + ) + + codec = codec_registry.get(component_quantization.quant_method) + bundle = codec.group( + descriptor, + source_weights, + component_quantization, + ) + for source_key in bundle.source_keys: + result.pop(source_key, None) + for record in bundle.values(): + if isinstance(record.storage, FloatWeight): + result[record.storage.source_key] = record.storage.value + continue + assert isinstance(record.storage, PackedWeight) + if "expert" in record.name: + raise NotImplementedError( + f"Packed expert weight {record.name!r} requires a " + "component-specific QMoE weight adapter." + ) + if component_quantization.tie_word_embeddings and any( + token in record.name for token in ("embed_tokens", "lm_head") + ): + raise NotImplementedError( + f"Tied packed table {record.name!r} requires a " + "component-specific tied-weight adapter." + ) + local_path = _local_weight_module_path(record.name, descriptor) + quantization = _effective_module_quantization( + component_quantization, + descriptor, + local_path, + ) + if quantization is None: + raise ValueError( + f"Packed checkpoint weight {record.name!r} targets a module " + f"excluded from component {component!r} quantization" + ) + result.update(codec.normalize(record, quantization)) + + canonical_keys = frozenset( + key + for component in component_names + for key in _canonical_component_parameter_keys( + module, + manifest[component], + ) + ) + remaining = next( + (key for key in result if _raw_qweight_key(key) and key not in canonical_keys), + None, + ) + if remaining is not None: + raise ValueError( + f"Packed checkpoint weight {remaining!r} was not routed to any " + "ModelPackage component" + ) + return result diff --git a/src/mobius/_component_quantization_test.py b/src/mobius/_component_quantization_test.py new file mode 100644 index 000000000..09d2b9f90 --- /dev/null +++ b/src/mobius/_component_quantization_test.py @@ -0,0 +1,268 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Tests for authoritative component quantization plans.""" + +from __future__ import annotations + +from typing import ClassVar + +import pytest +import torch +from onnxscript import nn + +from mobius._component_quantization import ( + configure_component_quantization, + normalize_component_quantized_weights, +) +from mobius._configs import ArchitectureConfig, QuantizationConfig +from mobius._model_package import ModelPackage +from mobius.components import ( + Linear, + QuantizedEmbedding, + QuantizedLinear, + make_quantized_linear_factory, +) +from mobius.tasks import ComponentSpec, ModelTask + + +class _DecoderLayer(nn.Module): + def __init__(self, linear_class: type[nn.Module]): + super().__init__() + self.q_proj = linear_class(64, 64, bias=False) + self.per_layer_input_gate = linear_class(64, 48, bias=False) + self.per_layer_projection = linear_class(64, 48, bias=False) + + +class _Backbone(nn.Module): + def __init__(self, linear_class: type[nn.Module]): + super().__init__() + self.layers = nn.ModuleList([_DecoderLayer(linear_class)]) + + +class _Decoder(nn.Module): + def __init__(self, linear_class: type[nn.Module]): + super().__init__() + self.model = _Backbone(linear_class) + + +class _Projection(nn.Module): + def __init__(self, linear_class: type[nn.Module] = Linear): + super().__init__() + self.proj = linear_class(64, 32, bias=False) + + +class _Composite(nn.Module): + HF_COMPONENT_SOURCES: ClassVar[dict[str, tuple[str, ...]]] = { + "decoder": ("model.language_model.layers", "lm_head"), + "audio_encoder": ("model.audio_tower",), + "embedding": ("model.language_model.embed_tokens",), + } + + def __init__(self): + super().__init__() + root_quantized = make_quantized_linear_factory(bits=4, block_size=16) + self.decoder = _Decoder(root_quantized) + self.audio_tower = _Projection() + self.embedding = _Projection(root_quantized) + + +class _CompositeTask(ModelTask): + model_roles: ClassVar[dict[str, str]] = { + "decoder": "decoder", + "audio_encoder": "encoder", + "embedding": "embedding", + } + components = ComponentSpec( + decoder="decoder", + audio_encoder="audio_tower", + embedding="embedding", + ) + + def build(self, module, config) -> ModelPackage: + raise NotImplementedError + + +def _config() -> ArchitectureConfig: + decoder = QuantizationConfig( + bits=4, + group_size=16, + quant_method="olive", + sym=True, + modules_to_not_convert=( + "lm_head", + r"re:.*\.per_layer_input_gate", + r"re:.*\.per_layer_projection", + ), + ) + return ArchitectureConfig( + quantization=decoder, + component_quantization={ + "decoder": decoder, + "audio_encoder": QuantizationConfig( + bits=8, + group_size=32, + quant_method="olive", + sym=True, + ), + }, + ) + + +def test_component_plan_applies_regex_exclusions_per_linear(): + module = _Composite() + + configure_component_quantization(module, _config(), _CompositeTask()) + + layer = module.decoder.model.layers[0] + assert isinstance(layer.q_proj, QuantizedLinear) + assert (layer.q_proj._bits, layer.q_proj._block_size) == (4, 16) + assert type(layer.per_layer_input_gate) is Linear + assert type(layer.per_layer_projection) is Linear + assert isinstance(module.audio_tower.proj, QuantizedLinear) + assert (module.audio_tower.proj._bits, module.audio_tower.proj._block_size) == ( + 8, + 32, + ) + # The mapping is authoritative: an omitted component stays float even + # though the top-level module was initially built from the decoder config. + assert type(module.embedding.proj) is Linear + + +def test_specialized_quantized_subclass_fails_instead_of_losing_semantics(): + class _SpecialQuantizedLinear(QuantizedLinear): + def forward(self, op, x): + return super().forward(op, x) + + module = _Projection(_SpecialQuantizedLinear) + config = ArchitectureConfig( + component_quantization={ + "model": QuantizationConfig( + bits=8, + group_size=32, + quant_method="olive", + ) + } + ) + + with pytest.raises(TypeError, match="specialized quantized module"): + configure_component_quantization(module, config, _SingleTask()) + + +def test_normalizes_weights_with_component_module_path_routing(): + module = _Composite() + config = _config() + task = _CompositeTask() + manifest = configure_component_quantization(module, config, task) + state_dict = { + "decoder.model.layers.0.q_proj.weight_qweight": torch.zeros(64, 32, dtype=torch.uint8), + "decoder.model.layers.0.q_proj.weight_scales": torch.ones(64, 4), + "decoder.model.layers.0.per_layer_input_gate.weight": torch.ones(48, 64), + "decoder.model.layers.0.per_layer_projection.weight": torch.ones(48, 64), + "audio_tower.proj.weight_qweight": torch.zeros(32, 64, dtype=torch.uint8), + "audio_tower.proj.weight_scales": torch.ones(32, 2), + "embedding.proj.weight": torch.ones(32, 64), + } + + result = normalize_component_quantized_weights( + state_dict, + module, + config, + ("decoder", "audio_encoder", "embedding"), + manifest=manifest, + task=task, + ) + + assert result["decoder.model.layers.0.q_proj.weight"].shape == (64, 4, 8) + assert result["audio_tower.proj.weight"].shape == (32, 2, 32) + assert result["decoder.model.layers.0.per_layer_input_gate.weight"].shape == ( + 48, + 64, + ) + + +def test_rejects_packed_weight_for_excluded_module(): + module = _Composite() + config = _config() + task = _CompositeTask() + manifest = configure_component_quantization(module, config, task) + state_dict = { + "decoder.model.layers.0.per_layer_input_gate.weight_qweight": torch.zeros( + 48, 32, dtype=torch.uint8 + ), + "decoder.model.layers.0.per_layer_input_gate.weight_scales": torch.ones(48, 4), + } + + with pytest.raises(ValueError, match="excluded"): + normalize_component_quantized_weights( + state_dict, + module, + config, + ("decoder", "audio_encoder", "embedding"), + manifest=manifest, + task=task, + ) + + +class _SingleTask(ModelTask): + model_roles: ClassVar[dict[str, str]] = {"model": "decoder"} + + def build(self, module, config) -> ModelPackage: + raise NotImplementedError + + +class _QuantizedEmbeddingModel(nn.Module): + def __init__(self): + super().__init__() + self.embed_tokens = QuantizedEmbedding( + 32, + 64, + bits=4, + block_size=16, + has_zero_point=False, + ) + self.proj = QuantizedLinear( + 64, + 32, + bits=4, + block_size=16, + has_zero_point=False, + ) + + +def test_canonical_quantized_embedding_is_not_treated_as_raw_sidecars(): + module = _QuantizedEmbeddingModel() + quantization = QuantizationConfig( + bits=4, + group_size=16, + quant_method="olive", + sym=True, + quantize_embeddings=True, + ) + config = ArchitectureConfig( + quantization=quantization, + component_quantization={"model": quantization}, + ) + task = _SingleTask() + manifest = configure_component_quantization(module, config, task) + state_dict = { + "embed_tokens.qweight": torch.zeros(32, 32, dtype=torch.uint8), + "embed_tokens.scales": torch.ones(32, 4), + "proj.weight": torch.zeros(32, 4, 8, dtype=torch.uint8), + "proj.scales": torch.ones(32, 4), + } + + result = normalize_component_quantized_weights( + state_dict, + module, + config, + ("model",), + manifest=manifest, + task=task, + ) + + assert set(result) == set(state_dict) + assert result["embed_tokens.qweight"] is state_dict["embed_tokens.qweight"] + assert result["embed_tokens.scales"] is state_dict["embed_tokens.scales"] + assert result["proj.weight"] is state_dict["proj.weight"] + assert result["proj.scales"] is state_dict["proj.scales"] diff --git a/src/mobius/_configs/__init__.py b/src/mobius/_configs/__init__.py index 8e3aea50d..17cd02f2d 100644 --- a/src/mobius/_configs/__init__.py +++ b/src/mobius/_configs/__init__.py @@ -90,7 +90,7 @@ _shallow_fields, _shared_expert_size, ) -from mobius._configs._quantization import QuantizationConfig +from mobius._configs._quantization import QuantizationConfig, QuantizationOverride from mobius._configs._sub_configs import ( AudioConfig, CodecDecoderConfig, @@ -153,6 +153,7 @@ "Plamo2Config", "QuantizationConfig", "Qwen4ExpConfig", + "QuantizationOverride", "Qwen35MtpConfig", "RoPEConfig", "Sam2Config", diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index e7c3cb954..ffabd439e 100644 --- a/src/mobius/_configs/_base.py +++ b/src/mobius/_configs/_base.py @@ -5,6 +5,7 @@ import dataclasses import math +from collections.abc import Mapping from typing import TYPE_CHECKING import onnx_ir as ir @@ -371,6 +372,104 @@ def _extract_audio_config(config, parent_config, model_type: str) -> dict: return _dispatch(config, parent_config, model_type) +_COMPONENT_QUANTIZATION_ALIASES = { + "text": "decoder", + "vision": "vision_encoder", + "audio": "audio_encoder", +} + + +def _get_config_value(config: object, name: str) -> object | None: + if isinstance(config, Mapping): + return config.get(name) + return getattr(config, name, None) + + +def _parse_component_quantization_mapping( + value: object, + *, + expert_dtype: object | None, +) -> dict[str, QuantizationConfig]: + if not isinstance(value, Mapping): + raise TypeError( + "component_quantization must be a mapping of component names to " + f"quantization configs, got {type(value).__name__}" + ) + + result: dict[str, QuantizationConfig] = {} + for raw_name, raw_config in value.items(): + if not isinstance(raw_name, str) or not raw_name: + raise ValueError("component_quantization keys must be non-empty strings") + name = _COMPONENT_QUANTIZATION_ALIASES.get(raw_name, raw_name) + quantization = QuantizationConfig.from_value( + raw_config, + expert_dtype=expert_dtype, + ) + if quantization is None: + continue + if name in result: + raise ValueError( + f"component_quantization declares component {name!r} more than once" + ) + result[name] = quantization + return result + + +def _extract_component_quantization( + config: object, + parent_config: object | None, + decoder_quantization: QuantizationConfig | None, +) -> dict[str, QuantizationConfig] | None: + """Extract authoritative explicit or nested component configurations.""" + sources = [] + for source in (parent_config, config): + if source is not None and all(id(source) != id(item) for item in sources): + sources.append(source) + + for source in sources: + declaration = _get_config_value(source, "component_quantization") + if declaration is None: + declaration = _get_config_value(source, "component_quantization_config") + if declaration is None: + root_quantization = _get_config_value(source, "quantization_config") + if hasattr(root_quantization, "to_dict"): + root_quantization = root_quantization.to_dict() + if isinstance(root_quantization, Mapping): + declaration = root_quantization.get("components") + if declaration is not None: + return _parse_component_quantization_mapping( + declaration, + expert_dtype=_get_config_value(source, "expert_dtype"), + ) + + composite = parent_config or config + nested: dict[str, QuantizationConfig] = {} + found_nested_declaration = False + for field_name, component_name in ( + ("vision_config", "vision_encoder"), + ("audio_config", "audio_encoder"), + ): + sub_config = _get_config_value(composite, field_name) + if sub_config is None: + continue + raw_quantization = _get_config_value(sub_config, "quantization_config") + if raw_quantization is None: + continue + found_nested_declaration = True + quantization = QuantizationConfig.from_value( + raw_quantization, + expert_dtype=_get_config_value(sub_config, "expert_dtype"), + ) + if quantization is not None: + nested[component_name] = quantization + + if not found_nested_declaration: + return None + if decoder_quantization is not None: + nested["decoder"] = dataclasses.replace(decoder_quantization) + return nested + + @dataclasses.dataclass class BaseModelConfig: """Base configuration shared by all model architectures. @@ -395,6 +494,9 @@ class BaseModelConfig: # Model dtype (from HF config dtype) dtype: ir.DataType = ir.DataType.FLOAT quantization: QuantizationConfig | None = None + # ``None`` keeps legacy model-wide behavior. A mapping is authoritative: + # omitted components remain floating point. + component_quantization: dict[str, QuantizationConfig] | None = None # HuggingFace identity and token metadata used by package persistence. model_type: str | None = None @@ -403,6 +505,27 @@ class BaseModelConfig: mask_token_id: int | None = None diffusion_shift_logits: bool = False + def quantization_for(self, component: str) -> QuantizationConfig | None: + """Return the effective quantization plan for one package component.""" + if self.component_quantization is None: + return self.quantization + candidates = { + "decoder": ("decoder", "model"), + "model": ("model", "decoder"), + "vision_encoder": ("vision_encoder", "vision"), + "vision": ("vision", "vision_encoder"), + "audio_encoder": ("audio_encoder", "audio"), + "audio": ("audio", "audio_encoder"), + }.get(component, (component,)) + return next( + ( + self.component_quantization[name] + for name in candidates + if name in self.component_quantization + ), + None, + ) + @dataclasses.dataclass class ArchitectureConfig(BaseModelConfig): @@ -1386,6 +1509,17 @@ def parse_quantization(source): quant = parse_quantization(config) if quant is None and parent_config is not None: quant = parse_quantization(parent_config) + component_quantization = _extract_component_quantization( + config, + parent_config, + quant, + ) + if component_quantization is not None: + options["component_quantization"] = component_quantization + quant = component_quantization.get( + "decoder", + component_quantization.get("model"), + ) if quant is not None: options["quantization"] = quant diff --git a/src/mobius/_configs/_quantization.py b/src/mobius/_configs/_quantization.py index 419b80a9b..33ea37e06 100644 --- a/src/mobius/_configs/_quantization.py +++ b/src/mobius/_configs/_quantization.py @@ -6,6 +6,50 @@ from __future__ import annotations import dataclasses +import re +from collections.abc import Mapping + + +def _compile_pattern(pattern: str) -> re.Pattern[str]: + try: + return re.compile(pattern) + except re.error as error: + raise ValueError(f"Invalid quantization regex {pattern!r}: {error}") from error + + +@dataclasses.dataclass(frozen=True) +class QuantizationOverride: + """Per-module affine layout override emitted by an upstream quantizer.""" + + bits: int | None = None + group_size: int | None = None + sym: bool | None = None + + @classmethod + def from_value(cls, value: object) -> QuantizationOverride: + """Parse one serialized module override.""" + if isinstance(value, cls): + return value + if hasattr(value, "to_dict"): + value = value.to_dict() + if not isinstance(value, Mapping): + raise TypeError( + f"quantization override must be a mapping, got {type(value).__name__}" + ) + return cls( + bits=value.get("bits"), + group_size=value.get("group_size"), + sym=value.get("sym", value.get("symmetric")), + ) + + def apply(self, config: QuantizationConfig) -> QuantizationConfig: + """Return *config* with this module override applied.""" + updates = { + name: value + for name, value in dataclasses.asdict(self).items() + if value is not None + } + return dataclasses.replace(config, **updates) @dataclasses.dataclass @@ -40,21 +84,30 @@ class QuantizationConfig: # RTN records this in its own config (``tie_word_embeddings``) and may clear # the model's top-level flag, so it is tracked here independently. tie_word_embeddings: bool = False + # HuggingFace full module names or ``re:``-prefixed full-match regexes that + # remain floating point inside this component. + modules_to_not_convert: tuple[str, ...] = () + # Literal HuggingFace module names or ``re:``-prefixed full-match regexes. + # Insertion order is significant: the first matching override wins. + overrides: dict[str, QuantizationOverride] = dataclasses.field(default_factory=dict) @classmethod - def from_transformers(cls, hf_config) -> QuantizationConfig | None: - """Parse ``quantization_config`` from a HuggingFace config. - - Returns ``None`` when no quantization is configured. - """ - qc = getattr(hf_config, "quantization_config", None) - if qc is None: + def from_value( + cls, + value: object, + *, + expert_dtype: object | None = None, + ) -> QuantizationConfig | None: + """Parse one serialized HuggingFace quantization configuration.""" + if value is None: return None - # qc can be a dict or a HF QuantizationConfig object - if hasattr(qc, "to_dict"): - qc = qc.to_dict() - if not isinstance(qc, dict): + if isinstance(value, cls): + return value + if hasattr(value, "to_dict"): + value = value.to_dict() + if not isinstance(value, Mapping): return None + qc = dict(value) method = qc.get("quant_method", "none") # NVIDIA ModelOpt NVFP4/FP8 checkpoints (e.g. quantized Qwen3.6) encode # weights as packed E2M1 (fp4) / float8 with block + global scales — a @@ -102,7 +155,8 @@ def from_transformers(cls, hf_config) -> QuantizationConfig | None: from mobius.integrations._block_quant import BlockQuantScheme scheme = BlockQuantScheme.from_quantization_config( - qc, expert_dtype=getattr(hf_config, "expert_dtype", None) + qc, + expert_dtype=expert_dtype, ) if scheme is not None: from mobius.integrations._block_quant import BlockQuantExportError @@ -128,13 +182,74 @@ def from_transformers(cls, hf_config) -> QuantizationConfig | None: # fp8 was already routed to the typed blocker above.) if method == "fp8": return None + raw_exclusions = qc.get("modules_to_not_convert") or () + if not isinstance(raw_exclusions, (list, tuple)): + raise TypeError( + "quantization_config.modules_to_not_convert must be a list or tuple" + ) + exclusions = tuple(str(pattern) for pattern in raw_exclusions) + raw_overrides = qc.get("overrides") or {} + if not isinstance(raw_overrides, Mapping): + raise TypeError("quantization_config.overrides must be a mapping") + overrides = { + str(pattern): QuantizationOverride.from_value(override) + for pattern, override in raw_overrides.items() + } + for pattern in (*exclusions, *overrides): + if pattern.startswith("re:"): + _compile_pattern(pattern[3:]) return cls( bits=qc.get("bits", 4), group_size=qc.get("group_size", 128), quant_method=method, sym=qc.get("sym", qc.get("symmetric", True)), - quantize_embeddings=bool(qc.get("embeds", False)), - quantize_lm_head=bool(qc.get("lm_head", False)), - quantize_vision=bool(qc.get("quantize_vision", False)), - tie_word_embeddings=bool(qc.get("tie_word_embeddings", False)), + float_zero_point=bool(qc.get("float_zero_point")), + quantize_embeddings=bool(qc.get("embeds")), + quantize_lm_head=bool(qc.get("lm_head")), + quantize_vision=bool(qc.get("quantize_vision")), + tie_word_embeddings=bool(qc.get("tie_word_embeddings")), + modules_to_not_convert=exclusions, + overrides=overrides, + ) + + @classmethod + def from_transformers(cls, hf_config) -> QuantizationConfig | None: + """Parse ``quantization_config`` from a HuggingFace config. + + Returns ``None`` when no quantization is configured. + """ + return cls.from_value( + getattr(hf_config, "quantization_config", None), + expert_dtype=getattr(hf_config, "expert_dtype", None), ) + + @staticmethod + def _matches_exclusion(pattern: str, module_name: str) -> bool: + if pattern.startswith("re:"): + return _compile_pattern(pattern[3:]).fullmatch(module_name) is not None + return pattern in module_name + + @staticmethod + def _matches_override(pattern: str, module_name: str) -> bool: + if pattern.startswith("re:"): + return _compile_pattern(pattern[3:]).fullmatch(module_name) is not None + return pattern == module_name + + def for_module( + self, + source_module_names: tuple[str, ...], + ) -> QuantizationConfig | None: + """Return this component's effective layout for one source module.""" + if any( + self._matches_exclusion(pattern, module_name) + for pattern in self.modules_to_not_convert + for module_name in source_module_names + ): + return None + for pattern, override in self.overrides.items(): + if any( + self._matches_override(pattern, module_name) + for module_name in source_module_names + ): + return override.apply(self) + return self diff --git a/src/mobius/_configs_test.py b/src/mobius/_configs_test.py index 24871b8c8..8678c556e 100644 --- a/src/mobius/_configs_test.py +++ b/src/mobius/_configs_test.py @@ -1058,6 +1058,77 @@ def test_from_transformers_olive_component_flags(self): assert qc.quantize_lm_head is True assert qc.quantize_vision is True + def test_component_plan_matches_exact_and_regex_module_rules(self): + qc = QuantizationConfig.from_value( + { + "quant_method": "olive", + "bits": 4, + "group_size": 32, + "modules_to_not_convert": [ + r"re:.*\.per_layer_input_gate", + ], + "overrides": { + "model.layers.0.q_proj": { + "bits": 8, + "group_size": 64, + } + }, + } + ) + + assert qc is not None + assert qc.for_module(("model.language_model.layers.0.per_layer_input_gate",)) is None + overridden = qc.for_module(("model.layers.0.q_proj",)) + assert overridden is not None + assert (overridden.bits, overridden.group_size) == (8, 64) + + def test_invalid_component_regex_fails_during_config_parse(self): + with pytest.raises(ValueError, match="Invalid quantization regex"): + QuantizationConfig.from_value( + { + "quant_method": "olive", + "modules_to_not_convert": ["re:("], + } + ) + + def test_architecture_config_parses_explicit_component_quantization(self): + text = SimpleNamespace( + model_type="llama", + hidden_size=64, + intermediate_size=128, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=4, + vocab_size=256, + hidden_act="silu", + max_position_embeddings=128, + ) + parent = SimpleNamespace( + model_type="composite", + component_quantization={ + "decoder": { + "quant_method": "olive", + "bits": 4, + "group_size": 32, + "modules_to_not_convert": [ + r"re:.*\.per_layer_projection", + ], + }, + "vision": { + "quant_method": "olive", + "bits": 8, + "group_size": 64, + }, + }, + ) + + config = ArchitectureConfig.from_transformers(text, parent_config=parent) + + assert config.component_quantization is not None + assert config.quantization_for("decoder").bits == 4 + assert config.quantization_for("vision_encoder").bits == 8 + assert config.quantization_for("decoder").modules_to_not_convert + def test_quantize_component_flags_default_false(self): qc = QuantizationConfig() assert qc.quantize_embeddings is False diff --git a/src/mobius/components/__init__.py b/src/mobius/components/__init__.py index d04b87e62..9bfc0f1b5 100644 --- a/src/mobius/components/__init__.py +++ b/src/mobius/components/__init__.py @@ -41,6 +41,7 @@ "GlmOcrVisionModel", "GatedShortConv", "ClippableLinear", + "ClippableQuantizedLinear", "GroupNorm", "GQAContext", "INT64_MAX", @@ -107,6 +108,7 @@ "get_activation", "initialize_rope", "make_quantized_linear_factory", + "make_clippable_quantized_linear_factory", "siglip2_naflex_attention_mask", ] @@ -260,10 +262,12 @@ ) from mobius.components._quantized_linear import ( BlockQuantizedLinear, + ClippableQuantizedLinear, NVFP4QuantizedLinear, QuantizedEmbedding, QuantizedLinear, TiedQuantizedLMHead, + make_clippable_quantized_linear_factory, make_quantized_linear_factory, ) from mobius.components._qwen3_asr_audio import ( diff --git a/src/mobius/components/_quantized_linear.py b/src/mobius/components/_quantized_linear.py index 4bcf2444c..94ab94014 100644 --- a/src/mobius/components/_quantized_linear.py +++ b/src/mobius/components/_quantized_linear.py @@ -253,6 +253,51 @@ def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value: return result +class ClippableQuantizedLinear(QuantizedLinear): + """Weight-quantized linear with learned input/output activation bounds.""" + + def __init__( + self, + in_features: int, + out_features: int, + bits: int = 4, + block_size: int = 32, + has_zero_point: bool = False, + zero_point_dtype: ir.DataType = ir.DataType.UINT8, + bias: bool = False, + ): + super().__init__( + in_features, + out_features, + bits, + block_size, + has_zero_point, + zero_point_dtype, + bias, + ) + self.input_min = nn.Parameter([]) + self.input_max = nn.Parameter([]) + self.output_min = nn.Parameter([]) + self.output_max = nn.Parameter([]) + + @staticmethod + def _clip( + op: OpBuilder, + x: ir.Value, + minimum: ir.Value, + maximum: ir.Value, + ) -> ir.Value: + x_f32 = op.Cast(x, to=ir.DataType.FLOAT) + minimum_f32 = op.Cast(minimum, to=ir.DataType.FLOAT) + maximum_f32 = op.Cast(maximum, to=ir.DataType.FLOAT) + return op.CastLike(op.Clip(x_f32, minimum_f32, maximum_f32), x) + + def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value: + x = self._clip(op, x, self.input_min, self.input_max) + result = super().forward(op, x) + return self._clip(op, result, self.output_min, self.output_max) + + class BlockQuantizedLinear(nn.Module): """Linear layer backed by native GGUF block quantization. @@ -547,3 +592,33 @@ def __init__( _Factory.__name__ = "QuantizedLinear" _Factory.__qualname__ = "QuantizedLinear" return _Factory + + +def make_clippable_quantized_linear_factory( + bits: int = 4, + block_size: int = 32, + has_zero_point: bool = False, + zero_point_dtype: ir.DataType = ir.DataType.UINT8, +) -> type[ClippableQuantizedLinear]: + """Create a Linear-compatible clipped MatMulNBits factory.""" + + class _Factory(ClippableQuantizedLinear): + def __init__( + self, + in_features: int, + out_features: int, + bias: bool = True, + ): + super().__init__( + in_features=in_features, + out_features=out_features, + bias=bias, + bits=bits, + block_size=block_size, + has_zero_point=has_zero_point, + zero_point_dtype=zero_point_dtype, + ) + + _Factory.__name__ = "ClippableQuantizedLinear" + _Factory.__qualname__ = "ClippableQuantizedLinear" + return _Factory diff --git a/src/mobius/components/_quantized_linear_test.py b/src/mobius/components/_quantized_linear_test.py index e1d4feb20..9f35c182c 100644 --- a/src/mobius/components/_quantized_linear_test.py +++ b/src/mobius/components/_quantized_linear_test.py @@ -19,6 +19,7 @@ ) from mobius.components._quantized_linear import ( BlockQuantizedLinear, + ClippableQuantizedLinear, NVFP4QuantizedLinear, QuantizedLinear, ) @@ -374,6 +375,35 @@ def test_rejects_runtime_unsupported_iq_format(self): BlockQuantizedLinear(IN_FEATURES, OUT_FEATURES, format="q4_k") +class TestClippableQuantizedLinear: + def test_keeps_clipping_parameters(self): + linear = ClippableQuantizedLinear(IN_FEATURES, OUT_FEATURES) + names = {name for name, _ in linear.named_parameters()} + assert { + "weight", + "scales", + "input_min", + "input_max", + "output_min", + "output_max", + } <= names + + def test_graph_clips_around_matmulnbits(self): + linear = ClippableQuantizedLinear( + IN_FEATURES, + OUT_FEATURES, + bits=8, + block_size=32, + ) + b, op, graph = create_test_builder() + x = create_test_input(b, "x", [1, 4, IN_FEATURES]) + result = linear(op, x) + b._adapt_outputs([result], "") + + assert count_op_type(graph, "MatMulNBits") == 1 + assert count_op_type(graph, "Clip") == 2 + + class TestMakeQuantizedLinearFactory: """Tests for the make_quantized_linear_factory closure.""" @@ -409,6 +439,22 @@ def test_factory_matches_linear_signature(self): assert instance._n == 64 assert instance.bias is None + def test_clippable_factory_uses_requested_layout(self): + from mobius.components._quantized_linear import ( + make_clippable_quantized_linear_factory, + ) + + factory = make_clippable_quantized_linear_factory( + bits=8, + block_size=32, + has_zero_point=True, + ) + linear = factory(IN_FEATURES, OUT_FEATURES, bias=False) + + assert isinstance(linear, ClippableQuantizedLinear) + assert linear.weight.shape == [OUT_FEATURES, 2, 32] + assert linear.zero_points is not None + class TestQuantizedEmbeddingInit: VOCAB = 64 diff --git a/src/mobius/components/_vision.py b/src/mobius/components/_vision.py index dd91a08b2..7b5d5d851 100644 --- a/src/mobius/components/_vision.py +++ b/src/mobius/components/_vision.py @@ -19,6 +19,7 @@ from onnxscript import OpBuilder, nn from mobius._configs import ArchitectureConfig +from mobius.components._common import Linear from mobius.components._mlp import FCMLP if TYPE_CHECKING: @@ -131,7 +132,7 @@ def forward( return self.out_proj(op, attn_output) -class _VisionLinear(nn.Module): +class _VisionLinear(Linear): """Linear layer with Transpose+MatMul using standard ONNX ops. Always includes bias. The ``bias`` kwarg is accepted for API @@ -139,14 +140,7 @@ class _VisionLinear(nn.Module): """ def __init__(self, in_features: int, out_features: int, bias: bool = True): - super().__init__() - self.weight = nn.Parameter([out_features, in_features]) - self.bias = nn.Parameter([out_features]) - - def forward(self, op: OpBuilder, x: ir.Value): - weight_t = op.Transpose(self.weight, perm=[1, 0]) - result = op.MatMul(x, weight_t) - return op.Add(result, self.bias) + super().__init__(in_features, out_features, bias=True) class VisionLayerNorm(nn.Module): diff --git a/src/mobius/integrations/transformers/_builder.py b/src/mobius/integrations/transformers/_builder.py index 459c60c0d..fe75adc49 100644 --- a/src/mobius/integrations/transformers/_builder.py +++ b/src/mobius/integrations/transformers/_builder.py @@ -13,6 +13,7 @@ from onnxscript import nn from mobius._builder import build_from_module, resolve_dtype +from mobius._component_quantization import normalize_component_quantized_weights from mobius._model_package import ModelPackage from mobius._registry import registry from mobius.integrations._weight_loading import ( @@ -56,6 +57,7 @@ def _strip_to_text_only(config: Any, model_type: str) -> Any: "boa_token_id", "vision", "audio", + "component_quantization", ): if name in field_names: overrides[name] = None @@ -327,6 +329,14 @@ def build_transformers_model( if task is None: task = _default_task_for_model(model_type) + from mobius.tasks import get_task + + resolved_task = get_task(task) + component_manifest = resolved_task.component_manifest( + module_class=module_class, + model_type=model_type, + hf_config=parent_config, + ) model_module = module_class(config) package = build_from_module( model_module, @@ -337,6 +347,7 @@ def build_transformers_model( fp8_kv_cache=fp8_kv_cache, kv_cache_scales=kv_cache_scales, prune_prefill_prefix=prune_prefill_prefix, + component_manifest=component_manifest, ) for name, model in package.items(): model.graph.name = f"{model_id}/{name}" @@ -401,6 +412,14 @@ def build_transformers_model( state_dict = _download_weights(model_id, revision=revision) if hasattr(model_module, "preprocess_weights"): state_dict = model_module.preprocess_weights(state_dict) + state_dict = normalize_component_quantized_weights( + state_dict, + model_module, + config, + package.keys(), + manifest=component_manifest, + task=resolved_task, + ) package.apply_weights( state_dict, prefix_map=getattr(model_module, "weight_prefix_map", None), diff --git a/src/mobius/integrations/transformers/_builder_test.py b/src/mobius/integrations/transformers/_builder_test.py index 02d5672c3..f54321ede 100644 --- a/src/mobius/integrations/transformers/_builder_test.py +++ b/src/mobius/integrations/transformers/_builder_test.py @@ -12,6 +12,7 @@ import pytest from onnxscript import nn +from mobius._configs import QuantizationConfig from mobius._model_package import ModelPackage from mobius._testing import make_config from mobius.integrations._block_quant import BlockQuantScheme @@ -265,6 +266,30 @@ def from_pretrained(model_id, **kwargs): assert calls == [("unsloth/Qwen3.8-Flash-Next-FP8", expected_kwargs)] +def test_strip_to_text_only_drops_component_quantization() -> None: + decoder = QuantizationConfig( + bits=4, + group_size=16, + quant_method="olive", + ) + config = make_config( + quantization=decoder, + component_quantization={ + "decoder": decoder, + "vision_encoder": QuantizationConfig( + bits=8, + group_size=32, + quant_method="olive", + ), + }, + ) + + stripped = transformers_builder._strip_to_text_only(config, "qwen2") + + assert stripped.component_quantization is None + assert stripped.quantization is decoder + + def test_transformers_build_uses_canonical_weight_loader(monkeypatch) -> None: hf_config = type("HFConfig", (), {"model_type": "qwen2"})() config = make_config(model_type="qwen2") diff --git a/src/mobius/models/gemma4_test.py b/src/mobius/models/gemma4_test.py index ff48f1c5a..8e8c90bd6 100644 --- a/src/mobius/models/gemma4_test.py +++ b/src/mobius/models/gemma4_test.py @@ -349,6 +349,39 @@ def test_quantized_decoder_does_not_quantize_vision_by_default(self): assert not any(node.op_type == "MatMulNBits" for node in graph) +def test_component_regex_keeps_per_layer_decoder_projections_float(): + from mobius._component_quantization import configure_component_quantization + from mobius.components import Linear, QuantizedLinear + from mobius.tasks._gemma4 import Gemma4Task + + quantization = QuantizationConfig( + bits=4, + group_size=16, + quant_method="olive", + sym=True, + modules_to_not_convert=( + "lm_head", + r"re:.*\.per_layer_input_gate", + r"re:.*\.per_layer_projection", + ), + ) + config = _tiny_gemma4_config( + enable_moe_block=False, + hidden_size_per_layer_input=16, + vocab_size_per_layer_input=256, + quantization=quantization, + component_quantization={"decoder": quantization}, + ) + module = Gemma4Model(config) + + configure_component_quantization(module, config, Gemma4Task()) + + layer = module.decoder.model.layers[0] + assert isinstance(layer.self_attn.q_proj, QuantizedLinear) + assert type(layer.per_layer_input_gate) is Linear + assert type(layer.per_layer_projection) is Linear + + class TestScaleFreeRMSNormOverflow: """V norm should handle FP16 overflow from squaring large values.""" diff --git a/src/mobius/models/mage_vl.py b/src/mobius/models/mage_vl.py index 522d4f6bd..ce09c29da 100644 --- a/src/mobius/models/mage_vl.py +++ b/src/mobius/models/mage_vl.py @@ -528,6 +528,7 @@ class MageVLEmbeddingModel(nn.Module): def __init__(self, config: ArchitectureConfig): super().__init__() + self.hidden_size = config.hidden_size self.embed_tokens = Embedding( config.vocab_size, config.hidden_size, @@ -557,7 +558,7 @@ def forward(self, op: OpBuilder, input_ids: Value, image_features: Value): indices = op.Reshape(flat_indices, op.Shape(input_ids)) flat_text = op.Reshape( text_embeddings, - op.Constant(value_ints=[-1, self.embed_tokens.weight.shape[1]]), + op.Constant(value_ints=[-1, self.hidden_size]), ) zero_feature = op.Mul( op.Slice(flat_text, starts=[0], ends=[1], axes=[0]), diff --git a/tests/build_graph_test.py b/tests/build_graph_test.py index 2686c30f6..b9bdd459d 100644 --- a/tests/build_graph_test.py +++ b/tests/build_graph_test.py @@ -16,6 +16,7 @@ from __future__ import annotations +import dataclasses import re import ml_dtypes @@ -57,6 +58,7 @@ AudioConfig, CodePredictorConfig, MMSConfig, + QuantizationConfig, SpeakerEncoderConfig, TTSConfig, VisionConfig, @@ -223,6 +225,34 @@ def _make_params( return params +def _with_component_quantization(config: ArchitectureConfig, task): + """Assign distinct tiny affine layouts to every materialized component.""" + layouts = ((4, 16), (8, 32), (2, 16)) + component_quantization = {} + for index, (component, role) in enumerate(task.model_roles.items()): + if role == "glue": + continue + if component == "audio_encoder" and config.audio is None: + continue + bits, group_size = layouts[index % len(layouts)] + component_quantization[component] = QuantizationConfig( + bits=bits, + group_size=group_size, + quant_method="olive", + sym=True, + quantize_embeddings=role == "embedding", + ) + decoder_quantization = component_quantization.get( + "decoder", + component_quantization.get("model"), + ) + return dataclasses.replace( + config, + quantization=decoder_quantization, + component_quantization=component_quantization, + ) + + # Configs imported from _test_configs — strip the is_representative flag # for use with pytest.parametrize. _MODEL_CONFIGS: list[tuple[str, dict]] = [(mt, ov) for mt, ov, _ in ALL_CAUSAL_LM_CONFIGS] @@ -530,6 +560,19 @@ def test_package_has_encoder_and_decoder(self, model_type: str, config_overrides dec_outputs = {out.name for out in pkg["decoder"].graph.outputs} assert "logits" in dec_outputs + def test_component_quantization_builds(self, model_type: str, config_overrides: dict): + config = _base_config(**config_overrides) + task = get_task(_default_task_for_model(model_type)) + config = _with_component_quantization(config, task) + + package = build_from_module( + registry.get(model_type)(config), + config, + task=task, + ) + + assert set(package) == {"encoder", "decoder"} + def test_onnx_checker_passes(self, model_type: str, config_overrides: dict): """Run the ONNX CheckerPass to catch attribute/shape/type errors.""" config = _base_config(**config_overrides) @@ -6835,6 +6878,35 @@ def test_package_builds(self, model_type: str, config_overrides: dict): pixel_values = next(i for i in vision.graph.inputs if i.name == "pixel_values") assert pixel_values.dtype == ir.DataType.FLOAT + def test_component_quantization_builds(self, model_type: str, config_overrides: dict): + """Every VL task accepts an independent affine layout per component.""" + config = _base_config(**config_overrides) + task = get_task(_default_task_for_model(model_type)) + config = _with_component_quantization(config, task) + + package = build_from_module( + registry.get(model_type)(config), + config, + task=task, + ) + + for component, model in package.items(): + quantization = config.component_quantization.get(component) + if quantization is None: + continue + quantized_nodes = [ + node + for node in model.graph + if node.op_type in {"MatMulNBits", "GatherBlockQuantized"} + ] + if not quantized_nodes: + assert not any(node.op_type == "MatMul" for node in model.graph), ( + f"{model_type}/{component} kept eligible MatMul projections float" + ) + for node in quantized_nodes: + assert node.attributes["bits"].as_int() == quantization.bits + assert node.attributes["block_size"].as_int() == quantization.group_size + def test_has_initializers(self, model_type: str, config_overrides: dict): """Verify all sub-models have non-empty initializers.""" config = _base_config(**config_overrides) @@ -6903,6 +6975,19 @@ def test_package_builds(self, model_type: str, config_overrides: dict): assert len(model.graph.inputs) > 0, f"{model_type}/{name} has no inputs" assert len(model.graph.outputs) > 0, f"{model_type}/{name} has no outputs" + def test_component_quantization_builds(self, model_type: str, config_overrides: dict): + config = _base_config(**config_overrides) + task = get_task(_default_task_for_model(model_type)) + config = _with_component_quantization(config, task) + + package = build_from_module( + registry.get(model_type)(config), + config, + task=task, + ) + + assert package + def test_has_initializers(self, model_type: str, config_overrides: dict): """Verify all sub-models have non-empty initializers.""" config = _base_config(**config_overrides) From e5ce365b091fe3b5a0e691f7e0438d8d3eb17460 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Thu, 27 Aug 2026 14:29:47 -0700 Subject: [PATCH 2/3] Format component loader after restack Apply Ruff formatting preserved across the manifest field rename rebase. Signed-off-by: Xiaoyu Zhang --- src/mobius/_component_quantization.py | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/src/mobius/_component_quantization.py b/src/mobius/_component_quantization.py index eee210cfe..5d7468c47 100644 --- a/src/mobius/_component_quantization.py +++ b/src/mobius/_component_quantization.py @@ -339,8 +339,7 @@ def configure_component_quantization( tuple( other.module_attribute_path for other in manifest.values() - if other.name != descriptor.name - and other.module_attribute_path + if other.name != descriptor.name and other.module_attribute_path ) if not descriptor.module_attribute_path else () @@ -372,9 +371,7 @@ def _canonical_component_parameter_keys( keys: set[str] = set() prefixes = { - prefix - for prefix in (descriptor.name, descriptor.module_attribute_path) - if prefix + prefix for prefix in (descriptor.name, descriptor.module_attribute_path) if prefix } for local_path, child in component_module.named_modules(): if not local_path: @@ -432,9 +429,7 @@ def owner(key: str) -> str | None: ] if not matches: root_components = [ - name - for name in component_names - if not manifest[name].module_attribute_path + name for name in component_names if not manifest[name].module_attribute_path ] return root_components[0] if len(root_components) == 1 else None max_length = max(length for length, _ in matches) From c6acc014f46bda0bce6753d43aa0e31d02dd9d21 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 4 Sep 2026 22:10:53 +0000 Subject: [PATCH 3/3] Fix post-merge config and quantization resolution for transformers builder Co-authored-by: xiaoyu-work <85524621+xiaoyu-work@users.noreply.github.com> --- src/mobius/_configs/_base.py | 227 ++++++++++++++++++ src/mobius/_configs/_quantization.py | 19 +- .../integrations/transformers/_builder.py | 1 - .../transformers/_builder_test.py | 12 - 4 files changed, 245 insertions(+), 14 deletions(-) diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index ffabd439e..9a89c145b 100644 --- a/src/mobius/_configs/_base.py +++ b/src/mobius/_configs/_base.py @@ -467,6 +467,7 @@ def _extract_component_quantization( return None if decoder_quantization is not None: nested["decoder"] = dataclasses.replace(decoder_quantization) + nested["embedding"] = dataclasses.replace(decoder_quantization) return nested @@ -526,6 +527,27 @@ def quantization_for(self, component: str) -> QuantizationConfig | None: None, ) + def quantization_for_source_paths( + self, + component: str, + source_paths: tuple[str, ...], + *, + ignored_source_names: tuple[str, ...] = ("lm_head", "embed_tokens"), + ) -> QuantizationConfig | None: + """Resolve module-level rules using one component's Hugging Face paths.""" + quantization = self.quantization_for(component) + if quantization is None or not quantization.has_module_plan: + return quantization + effective_paths = tuple( + path + for path in source_paths + if path.rsplit(".", 1)[-1] not in ignored_source_names + ) + return quantization.for_source_paths( + effective_paths or source_paths, + component=component, + ) + @dataclasses.dataclass class ArchitectureConfig(BaseModelConfig): @@ -1682,6 +1704,31 @@ class CausalLMConfig(ArchitectureConfig): """ +@dataclasses.dataclass +class GrokGGUFConfig(CausalLMConfig): + """GGUF-only Grok graph settings loaded by the pinned llama.cpp implementation.""" + + embedding_scale: float = 78.38367176906169 + attention_output_scale: float = 0.08838834764831845 + logit_output_scale: float = 0.5773502691896257 + attn_logit_softcapping: float = 30.0 + router_logit_softcapping: float = 30.0 + final_logit_softcapping: float = 0.0 + attention_temperature_length: int = 0 + has_dense_ffn: bool = False + has_gated_dense_ffn: bool = False + has_gated_experts: bool = True + + +@dataclasses.dataclass +class GroveMoEGGUFConfig(CausalLMConfig): + """GGUF-only GroveMoE chunk-expert routing settings.""" + + chunk_expert_intermediate_size: int = DEFAULT_INT + experts_per_group: int = DEFAULT_INT + expert_group_scale: float = 0.05 + + @dataclasses.dataclass class Qwen4ExpConfig(CausalLMConfig): """Exact configuration for experimental Qwen4/Qwen3.8 Flash-Next.""" @@ -2758,10 +2805,12 @@ class Eagle3Config(CausalLMConfig): eagle_aux_hidden_state_layer_ids: list[int] | None = None target_layer_ids: list[int] | None = None use_target_lm_head: bool = False + use_draft_token_embedding: bool = False @classmethod def from_transformers(cls, config, parent_config=None) -> Eagle3Config: layer_cfg = getattr(config, "transformer_layer_config", None) + is_speculators_format = layer_cfg is not None if layer_cfg is not None: # speculators format: arch config nested under transformer_layer_config. if isinstance(layer_cfg, dict): @@ -2782,6 +2831,7 @@ def from_transformers(cls, config, parent_config=None) -> Eagle3Config: eagle_aux_hidden_state_layer_ids=getattr( config, "eagle_aux_hidden_state_layer_ids", None ), + use_draft_token_embedding=is_speculators_format, ) @@ -4434,6 +4484,16 @@ class SpeechToTextConfig(ArchitectureConfig): decoder_start_token_id: int | None = None layer_norm_eps: float = 1e-5 + @property + def encoder_output_size(self) -> int: + """Channel width of ``encoder_hidden_states``. + + Defaults to the decoder width because most speech encoder-decoders share + one model dimension. Architectures whose encoder is narrower or wider + (e.g. Moonshine Streaming with a projection adapter) override this. + """ + return self.hidden_size + @dataclasses.dataclass class WhisperConfig(SpeechToTextConfig): @@ -4580,6 +4640,173 @@ def from_transformers(cls, config, parent_config=None) -> MoonshineConfig: return cls(**options) +def _sub_config_get(config, name: str, default): + """Read ``name`` from a sub-config that may be an object or a plain dict.""" + if isinstance(config, dict): + value = config.get(name, default) + else: + value = getattr(config, name, default) + return default if value is None else value + + +@dataclasses.dataclass +class MoonshineStreamingConfig(SpeechToTextConfig): + """Configuration for Moonshine Streaming raw-waveform encoder-decoder ASR. + + Moonshine Streaming replaces the offline Moonshine convolutional stem with a + fixed-length framing front end (``frame_ms`` frames of raw samples), per-frame + CMVN, learned asinh compression, and two *causal* strided convolutions. Its + encoder carries no rotary embedding; position information comes from the + causal stem plus per-layer asymmetric ``(left, right)`` sliding windows, which + bound the streaming lookahead. The decoder adds an absolute learned position + table (``pos_emb``) to the encoder output before cross-attention. + """ + + encoder_input_name: str = "input_values" + encoder_input_channels: int | None = None + encoder_uses_attention_mask: bool = True + decoder_uses_encoder_attention_mask: bool = True + encoder_hidden_size: int = DEFAULT_INT + encoder_intermediate_size: int = DEFAULT_INT + encoder_num_hidden_layers: int = DEFAULT_INT + encoder_num_attention_heads: int = DEFAULT_INT + encoder_num_key_value_heads: int = DEFAULT_INT + encoder_head_dim: int = DEFAULT_INT + encoder_hidden_act: str = "gelu" + decoder_hidden_act: str = "silu" + #: Q/K/V and output projection bias of the encoder attention. Upstream + #: gates all four encoder projections on the encoder sub-config's + #: ``attention_bias`` (the decoder's output projection stays bias-free). + encoder_attention_bias: bool = False + #: Per-layer ``(left_window, right_window)`` attention spans of the encoder. + #: ``left`` counts the query position itself; ``right`` is the strict + #: lookahead. ``right == 0`` means the layer is fully causal. + encoder_sliding_windows: tuple[tuple[int, int], ...] = ((16, 4),) + encoder_sample_rate: int = 16_000 + encoder_frame_ms: float = 5.0 + #: Epsilon of the per-frame cepstral mean/variance normalisation. + encoder_cmvn_eps: float = 1e-6 + + def __post_init__(self): + # A tuple keeps the config hashable and prevents accidental mutation of + # the per-layer window schedule shared by every encoder layer. + self.encoder_sliding_windows = tuple( + (int(left), int(right)) for left, right in self.encoder_sliding_windows + ) + if self.encoder_hidden_size == DEFAULT_INT: + self.encoder_hidden_size = self.hidden_size + if self.encoder_intermediate_size == DEFAULT_INT: + self.encoder_intermediate_size = self.intermediate_size + if self.encoder_num_attention_heads == DEFAULT_INT: + self.encoder_num_attention_heads = self.num_attention_heads + if self.encoder_num_key_value_heads == DEFAULT_INT: + self.encoder_num_key_value_heads = self.encoder_num_attention_heads + if self.encoder_num_hidden_layers == DEFAULT_INT: + self.encoder_num_hidden_layers = len(self.encoder_sliding_windows) + if self.encoder_head_dim == DEFAULT_INT: + self.encoder_head_dim = ( + self.encoder_hidden_size // self.encoder_num_attention_heads + ) + if len(self.encoder_sliding_windows) != self.encoder_num_hidden_layers: + raise ValueError( + "MoonshineStreamingConfig: encoder_sliding_windows has " + f"{len(self.encoder_sliding_windows)} entries but the encoder has " + f"{self.encoder_num_hidden_layers} layers." + ) + + @property + def encoder_output_size(self) -> int: + """Encoder width; the decoder projects it when it differs from its own.""" + return self.encoder_hidden_size + + @property + def frame_length(self) -> int: + """Raw samples per encoder frame (``sample_rate * frame_ms / 1000``). + + Upstream rounds to the nearest integer, so the audio length fed to the + encoder must be a multiple of this value; the processor's + ``pad_to_multiple_of`` enforces that. + """ + return round(self.encoder_sample_rate * self.encoder_frame_ms / 1000.0) + + @classmethod + def from_transformers(cls, config, parent_config=None) -> MoonshineStreamingConfig: + if config.model_type != "moonshine_streaming": + raise ValueError( + "MoonshineStreamingConfig expects model_type='moonshine_streaming', " + f"got '{config.model_type}'" + ) + + hidden_size = config.hidden_size + decoder_heads = config.num_attention_heads + # ``encoder_config`` is a nested MoonshineStreamingEncoderConfig on a + # trusted config object and a plain dict when the JSON is read directly. + encoder_config = getattr(config, "encoder_config", None) or {} + encoder_hidden_size = _sub_config_get(encoder_config, "hidden_size", hidden_size) + encoder_heads = _sub_config_get(encoder_config, "num_attention_heads", decoder_heads) + rope_parameters = getattr(config, "rope_parameters", None) or getattr( + config, "rope_scaling", None + ) + rope_parameters = rope_parameters or {} + windows = _sub_config_get( + encoder_config, + "sliding_windows", + ((16, 4), (16, 4), (16, 0), (16, 0), (16, 4), (16, 4)), + ) + + options = dict( + vocab_size=config.vocab_size, + hidden_size=hidden_size, + intermediate_size=config.intermediate_size, + num_hidden_layers=config.num_hidden_layers, + num_attention_heads=decoder_heads, + num_key_value_heads=getattr(config, "num_key_value_heads", decoder_heads), + head_dim=getattr(config, "head_dim", None) or hidden_size // decoder_heads, + hidden_act=getattr(config, "hidden_act", "silu"), + pad_token_id=getattr(config, "pad_token_id", 0), + tie_word_embeddings=getattr(config, "tie_word_embeddings", False), + attn_qkv_bias=getattr(config, "attention_bias", False), + attn_o_bias=False, + max_position_embeddings=getattr(config, "max_position_embeddings", 4096), + rope_type=rope_parameters.get("rope_type", "default"), + rope_theta=rope_parameters.get("rope_theta", 10_000.0), + rope_scaling=rope_parameters or None, + partial_rotary_factor=rope_parameters.get("partial_rotary_factor", 1.0), + rope_interleave=True, + mlp_bias=True, + encoder_hidden_size=encoder_hidden_size, + encoder_intermediate_size=_sub_config_get( + encoder_config, "intermediate_size", config.intermediate_size + ), + encoder_num_hidden_layers=_sub_config_get( + encoder_config, "num_hidden_layers", config.num_hidden_layers + ), + encoder_num_attention_heads=encoder_heads, + encoder_num_key_value_heads=_sub_config_get( + encoder_config, "num_key_value_heads", encoder_heads + ), + encoder_head_dim=_sub_config_get(encoder_config, "head_dim", None) + or encoder_hidden_size // encoder_heads, + encoder_hidden_act=_sub_config_get(encoder_config, "hidden_act", "gelu"), + encoder_attention_bias=bool( + _sub_config_get(encoder_config, "attention_bias", False) + ), + decoder_hidden_act=getattr(config, "hidden_act", "silu"), + encoder_sliding_windows=tuple(tuple(window) for window in windows), + encoder_sample_rate=_sub_config_get(encoder_config, "sample_rate", 16_000), + encoder_frame_ms=_sub_config_get(encoder_config, "frame_ms", 5.0), + decoder_start_token_id=getattr(config, "decoder_start_token_id", 1), + layer_norm_eps=getattr(config, "layer_norm_eps", 1e-5), + model_type="moonshine_streaming", + bos_token_id=getattr(config, "bos_token_id", 1), + eos_token_id=getattr(config, "eos_token_id", 2), + ) + resolved = _resolve_dtype(config) + if resolved is not None: + options["dtype"] = resolved + return cls(**options) + + def _conv_widths(config, defaults, hidden_size: int) -> tuple[int, ...]: """Per-layer channel widths of a wav2vec2-family convolutional feature encoder. diff --git a/src/mobius/_configs/_quantization.py b/src/mobius/_configs/_quantization.py index 33ea37e06..0d8b8166e 100644 --- a/src/mobius/_configs/_quantization.py +++ b/src/mobius/_configs/_quantization.py @@ -223,6 +223,11 @@ def from_transformers(cls, hf_config) -> QuantizationConfig | None: expert_dtype=getattr(hf_config, "expert_dtype", None), ) + @property + def has_module_plan(self) -> bool: + """Whether module-selection rules or overrides are present.""" + return bool(self.modules_to_not_convert) or bool(self.overrides) + @staticmethod def _matches_exclusion(pattern: str, module_name: str) -> bool: if pattern.startswith("re:"): @@ -233,7 +238,7 @@ def _matches_exclusion(pattern: str, module_name: str) -> bool: def _matches_override(pattern: str, module_name: str) -> bool: if pattern.startswith("re:"): return _compile_pattern(pattern[3:]).fullmatch(module_name) is not None - return pattern == module_name + return pattern == module_name or module_name.startswith(f"{pattern}.") or pattern in module_name def for_module( self, @@ -253,3 +258,15 @@ def for_module( ): return override.apply(self) return self + + def for_source_paths( + self, + source_paths: tuple[str, ...], + *, + component: str = "decoder", + ) -> QuantizationConfig | None: + """Resolve module-level rules for a set of source paths.""" + resolved = self.for_module(source_paths) + if resolved is None: + return None + return dataclasses.replace(resolved, overrides={}, modules_to_not_convert=()) diff --git a/src/mobius/integrations/transformers/_builder.py b/src/mobius/integrations/transformers/_builder.py index 36cb1a95a..005547f1d 100644 --- a/src/mobius/integrations/transformers/_builder.py +++ b/src/mobius/integrations/transformers/_builder.py @@ -604,7 +604,6 @@ def build_transformers_model( manifest=component_manifest, task=resolved_task, ) - ) package.apply_weights( state_dict, prefix_map=getattr(model_module, "weight_prefix_map", None), diff --git a/src/mobius/integrations/transformers/_builder_test.py b/src/mobius/integrations/transformers/_builder_test.py index 392a72c25..831fc4e9a 100644 --- a/src/mobius/integrations/transformers/_builder_test.py +++ b/src/mobius/integrations/transformers/_builder_test.py @@ -305,7 +305,6 @@ def test_strip_to_text_only_drops_component_quantization() -> None: quant_method="olive", ) config = make_config( -<<<<<<< HEAD quantization=decoder, component_quantization={ "decoder": decoder, @@ -315,12 +314,6 @@ def test_strip_to_text_only_drops_component_quantization() -> None: quant_method="olive", ), }, -======= - component_quantization={ - "decoder": decoder, - "vision_encoder": decoder, - } ->>>>>>> origin/refactor/typed-weight-pipeline ) stripped = transformers_builder._strip_to_text_only(config, "qwen2") @@ -329,8 +322,6 @@ def test_strip_to_text_only_drops_component_quantization() -> None: assert stripped.quantization is decoder -<<<<<<< HEAD -======= def test_strip_to_text_only_resolves_decoder_module_plan() -> None: decoder = QuantizationConfig( bits=4, @@ -356,9 +347,6 @@ def test_strip_to_text_only_resolves_decoder_module_plan() -> None: assert stripped.quantization is not None assert (stripped.quantization.bits, stripped.quantization.group_size) == (8, 32) assert stripped.quantization.overrides == {} - - ->>>>>>> origin/refactor/typed-weight-pipeline def test_transformers_build_uses_canonical_weight_loader(monkeypatch) -> None: hf_config = type("HFConfig", (), {"model_type": "qwen2"})() config = make_config(model_type="qwen2")