diff --git a/CHANGELOG.md b/CHANGELOG.md index 973e57243..3a56b6a2e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,25 +7,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] -### Independent quantization for multi-component packages +### Per-component quantized checkpoint loading #### Added -- HuggingFace composite checkpoints can declare a `component_quantization` - mapping (or `quantization_config.components`) whose keys match - `ModelPackage` component names such as `decoder`, `encoder`, - `vision_encoder`, `audio_encoder`, and `embedding`. Nested - `vision_config.quantization_config` and `audio_config.quantization_config` - values are also recognized. -- `build_from_module` now configures every component independently. Existing - quantized decoder modules are retargeted to the component's bit width and - group size, float encoder/vision/audio projections are converted to - `MatMulNBits`, quantized embeddings use `GatherBlockQuantized`, and components - omitted from the mapping remain floating point. -- Olive mixed-precision component-wide `modules_to_not_convert` and `overrides` - are collapsed into component layouts. Partial module rules and genuinely - mixed layouts inside one ONNX component fail with an actionable error instead - of loading packed weights with the wrong configuration. +- 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 diff --git a/src/mobius/__main__.py b/src/mobius/__main__.py index 8a5a86a78..01cce5004 100644 --- a/src/mobius/__main__.py +++ b/src/mobius/__main__.py @@ -521,6 +521,15 @@ 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) + from mobius.tasks import get_task + from mobius._component_quantization import attach_hf_component_sources + + 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) attach_hf_component_sources( model_module, @@ -535,6 +544,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}" @@ -566,15 +576,16 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask: if hasattr(model_module, "preprocess_weights"): state_dict = model_module.preprocess_weights(state_dict) from mobius._component_quantization import ( - preprocess_component_quantized_state_dict, + normalize_component_quantized_weights, ) - state_dict = preprocess_component_quantized_state_dict( + state_dict = normalize_component_quantized_weights( state_dict, model_module, config, - task, pkg.keys(), + manifest=component_manifest, + task=resolved_task, ) pkg.apply_weights(state_dict) else: diff --git a/src/mobius/_builder.py b/src/mobius/_builder.py index 1abf62d62..bb2d3ea82 100644 --- a/src/mobius/_builder.py +++ b/src/mobius/_builder.py @@ -25,6 +25,7 @@ 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 @@ -137,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. @@ -162,8 +164,12 @@ def build_from_module( if prune_prefill_prefix: task = _enable_prefill_prefix_pruning_task(task) resolved_task = get_task(task) - component_manifest = resolved_task.component_manifest() - configure_component_quantization(module, config, resolved_task) + 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): diff --git a/src/mobius/_component_quantization.py b/src/mobius/_component_quantization.py index 5a44cf63e..40ab5c471 100644 --- a/src/mobius/_component_quantization.py +++ b/src/mobius/_component_quantization.py @@ -1,18 +1,25 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. -"""Apply independent weight-quantization layouts to package components.""" +"""Configure and load independently quantized model-package components.""" from __future__ import annotations -from collections.abc import Iterable +__all__ = [ + "attach_hf_component_sources", + "configure_component_quantization", + "normalize_component_quantized_weights", + "preprocess_component_quantized_state_dict", +] + +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._weight_utils import preprocess_quantized_weights from mobius.components import ( ClippableLinear, ClippableQuantizedLinear, @@ -24,25 +31,11 @@ make_quantized_linear_factory, ) from mobius.tasks import ModelTask, get_task +from mobius.weights import FloatWeight, PackedWeight, codec_registry -_AFFINE_QUANT_METHODS = frozenset({"olive", "gptq", "awq"}) -_COMPONENT_ATTRIBUTE_ALIASES = { - "vision": "vision_encoder", - "audio": "audio_encoder", - "speech": "speech_encoder", -} -_TOKEN_EMBEDDING_NAMES = frozenset( - { - "embed_in", - "embed_tokens", - "shared", - "word_embeddings", - "wte", - } -) -_KNOWN_COMPONENT_NAMES = frozenset( +_AFFINE_METHODS = frozenset({"olive", "gptq", "awq"}) +_KNOWN_SPLIT_COMPONENTS = frozenset( { - "model", "decoder", "encoder", "vision", @@ -50,46 +43,14 @@ "audio", "audio_encoder", "embedding", + "model", } ) -def attach_hf_component_sources( - module: nn.Module, - *, - model_type: str, - hf_config: object, -) -> None: - """Attach the runtime HF component map selected for this concrete model.""" - resolver = getattr(type(module), "get_hf_component_sources", None) - if resolver is not None: - source_map = resolver(model_type=model_type, hf_config=hf_config) - else: - source_map = getattr(type(module), "HF_COMPONENT_SOURCES", {}) - module._hf_component_sources = { - component: tuple(paths) for component, paths in source_map.items() - } - - -def _component_source_map(module: nn.Module) -> dict[str, tuple[str, ...]]: - return getattr( - module, - "_hf_component_sources", - getattr(type(module), "HF_COMPONENT_SOURCES", {}), - ) - - -def _component_output_head_paths(module: nn.Module, component: str) -> tuple[str, ...]: - mapping = getattr(type(module), "COMPONENT_OUTPUT_HEADS", {}) - aliases = { - "decoder": ("decoder", "model"), - "model": ("model", "decoder"), - }.get(component, (component,)) - declared = next((tuple(mapping[name]) for name in aliases if name in mapping), ()) - return tuple(dict.fromkeys(("lm_head", *declared))) - - -def _resolve_path(root: nn.Module, path: str) -> nn.Module | None: +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): @@ -98,60 +59,36 @@ def _resolve_path(root: nn.Module, path: str) -> nn.Module | None: return current if isinstance(current, nn.Module) else None -def _component_module_paths( - module: nn.Module, - task: ModelTask, -) -> dict[str, str]: - """Resolve package component names to ONNXScript module paths.""" - paths: dict[str, str] = {} - if task.components is not None: - paths.update( - { - component: path - for component, path in task.components.items() - if _resolve_path(module, path) is not None - } - ) - - for component in task.model_roles: - if component in paths: - continue - if component == "model": - paths[component] = "" - continue - candidates = ( - component, - _COMPONENT_ATTRIBUTE_ALIASES.get(component, component), - ) - for candidate in candidates: - if _resolve_path(module, candidate) is not None: - paths[component] = candidate - break - return paths - - -def _component_module(module: nn.Module, path: str) -> nn.Module: - if not path: - return module - resolved = _resolve_path(module, path) - if resolved is None: - raise ValueError(f"Cannot resolve component module path {path!r}") - return resolved - - -def _replace_child_module(root: nn.Module, path: str, replacement: nn.Module) -> None: - """Replace a named ONNXScript child while retaining its graph name.""" +def _replace_child(root: nn.Module, path: str, replacement: nn.Module) -> None: if not path: - raise ValueError("Cannot replace the root component module") + raise ValueError("Cannot replace a component's root module") parts = path.split(".") parent: object = root for part in parts[:-1]: parent = getattr(parent, part) - child_name = parts[-1] - old = getattr(parent, child_name) + old = getattr(parent, parts[-1]) if hasattr(replacement, "_set_name") and hasattr(old, "name"): replacement._set_name(old.name) - setattr(parent, child_name, replacement) + 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( @@ -180,305 +117,142 @@ def _clippable_linear_factory( ) -def _effective_component_quantization( - module: nn.Module, - config: BaseModelConfig, - component: str, -) -> QuantizationConfig | None: - resolver = getattr(config, "quantization_for", None) - if resolver is not None: - quantization = resolver(component) - else: - component_quantization = getattr(config, "component_quantization", None) - quantization = ( - component_quantization.get(component) - if component_quantization is not None - else getattr(config, "quantization", None) - ) - if quantization is None or quantization.quant_method == "none": - return None - if not quantization.has_module_plan: - return quantization - - source_map = _component_source_map(module) - source_paths = tuple(source_map.get(component, ())) - if not source_paths: - raise ValueError( - f"Component {component!r} carries module-level quantization rules, " - f"but {type(module).__name__} declares no HF_COMPONENT_SOURCES entry " - "from which Mobius can derive a uniform component layout." - ) - return config.quantization_for_source_paths( - component, - source_paths, - ignored_source_names=( - *_component_output_head_paths(module, component), - "embed_tokens", - ), - ) - - def _float_linear(module: QuantizedLinear) -> Linear: return Linear(module._k, module._n, bias=module.bias is not None) -class _ScaledEmbedding(Embedding): - """Float token embedding that preserves a model's post-gather scale.""" - - def __init__( - self, - num_embeddings: int, - embedding_dim: int, - padding_idx: int | None, - *, - embed_scale: float, - ): - super().__init__(num_embeddings, embedding_dim, padding_idx) - self.embed_scale = embed_scale - - def forward(self, op, input_ids): - return op.Mul(super().forward(op, input_ids), self.embed_scale) - - -class _ScaledQuantizedEmbedding(QuantizedEmbedding): - """Quantized token embedding that preserves a model's post-gather scale.""" - - def __init__( - self, - num_embeddings: int, - embedding_dim: int, - *, - bits: int, - block_size: int, - has_zero_point: bool, - padding_idx: int | None, - embed_scale: float, - ): - super().__init__( - num_embeddings, - embedding_dim, - bits=bits, - block_size=block_size, - has_zero_point=has_zero_point, - padding_idx=padding_idx, - ) - self.embed_scale = embed_scale - - def forward(self, op, input_ids): - return op.Mul(super().forward(op, input_ids), self.embed_scale) - - def _float_embedding(module: QuantizedEmbedding) -> Embedding: - args = ( + return Embedding( int(module.qweight.shape[0]), module._embedding_dim, module.padding_idx, ) - embed_scale = getattr(module, "embed_scale", None) - if embed_scale is not None: - return _ScaledEmbedding(*args, embed_scale=float(embed_scale)) - if type(module).forward is not QuantizedEmbedding.forward: - raise TypeError( - "Component plan cannot convert specialized quantized embedding " - f"{type(module).__name__} to a plain embedding without dropping " - "its forward semantics." - ) - return Embedding(*args) -def _quantized_embedding( - module: Embedding | QuantizedEmbedding, - quantization: QuantizationConfig, -) -> QuantizedEmbedding: - if isinstance(module, QuantizedEmbedding): - num_embeddings = int(module.qweight.shape[0]) - embedding_dim = module._embedding_dim - else: - num_embeddings, embedding_dim = (int(dim) for dim in module.weight.shape) - embed_scale = getattr(module, "embed_scale", None) - if embed_scale is not None: - return _ScaledQuantizedEmbedding( - num_embeddings, - embedding_dim, - bits=quantization.bits, - block_size=quantization.group_size, - has_zero_point=not quantization.sym, - padding_idx=module.padding_idx, - embed_scale=float(embed_scale), - ) - if ( - isinstance(module, QuantizedEmbedding) - and type(module).forward is not QuantizedEmbedding.forward - ) or (isinstance(module, Embedding) and type(module).forward is not Embedding.forward): - raise TypeError( - "Component plan cannot quantize specialized embedding " - f"{type(module).__name__} without dropping its forward semantics." - ) - return QuantizedEmbedding( - num_embeddings, - embedding_dim, - bits=quantization.bits, - block_size=quantization.group_size, - has_zero_point=not quantization.sym, - padding_idx=module.padding_idx, - ) - - -def _embedding_layout_matches( - module: QuantizedEmbedding, - quantization: QuantizationConfig, -) -> bool: - return ( - quantization.quantize_embeddings - and module._bits == quantization.bits - and module._block_size == quantization.group_size - and (module.zero_points is None) is quantization.sym - ) - - -def _excluded_from_component_quantization( - root: nn.Module, - path: str, - quantization: QuantizationConfig, -) -> bool: - """Return whether a model-declared subtree stays float for this method.""" - parts = path.split(".") - for end in range(1, len(parts) + 1): - module = _resolve_path(root, ".".join(parts[:end])) - methods = getattr(module, "component_quantization_excluded_methods", ()) - if quantization.quant_method in methods: - return True - return False - - -def _component_token_embedding_keys( - module: nn.Module, - component: str, - component_path: str, -) -> tuple[str, ...]: - """Return canonical float keys for token tables owned by one component.""" - component_module = _component_module(module, component_path) - prefixes = {component_path, component} if component_path else {""} - keys: set[str] = set() - for name, child in component_module.named_modules(): - if ( - name - and isinstance(child, (Embedding, QuantizedEmbedding)) - and name.rsplit(".", 1)[-1] in _TOKEN_EMBEDDING_NAMES - ): - for prefix in prefixes: - path = f"{prefix}.{name}" if prefix else name - keys.add(f"{path}.weight") - return tuple(sorted(keys)) - - -def _packed_qweight_for(key: str, float_key: str) -> bool: - owner = float_key.removesuffix(".weight") - return key in {f"{float_key}_qweight", f"{owner}.qweight"} +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, - quantization: QuantizationConfig | None, + component_quantization: QuantizationConfig | None, *, - output_head_paths: tuple[str, ...], + owned_by_other_components: tuple[str, ...] = (), ) -> None: - """Rewrite float/quantized projection scaffolding for one component.""" - named_modules = list(component_module.named_modules()) - - linear_factory = ( - _linear_factory(config, quantization) if quantization is not None else None - ) - clippable_factory = ( - _clippable_linear_factory(config, quantization) if quantization is not None else None - ) replacements: list[tuple[str, nn.Module]] = [] - - for name, child in named_modules: - if not name: + for local_path, child in list(component_module.named_modules()): + if not local_path: continue - is_lm_head = ( - name == "lm_head" or name.endswith(".lm_head") or name in output_head_paths - ) - excluded = quantization is not None and _excluded_from_component_quantization( - component_module, - name, - quantization, + 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 ( - quantization is None - or excluded - or (is_lm_head and not quantization.quantize_lm_head) - ): - replacement: nn.Module = ClippableLinear( + 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: - assert clippable_factory is not None - replacement = clippable_factory( + replacement = _clippable_linear_factory(config, quantization)( child._k, child._n, bias=child.bias is not None, ) - replacements.append((name, replacement)) + replacements.append((local_path, replacement)) continue if isinstance(child, QuantizedLinear): - if ( - quantization is None - or excluded - or (is_lm_head and not quantization.quantize_lm_head) - ): - replacement = _float_linear(child) - else: - assert linear_factory is not None - replacement = linear_factory( + 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((name, replacement)) + ) + 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((name, _float_embedding(child))) - elif not _embedding_layout_matches(child, quantization): - replacements.append((name, _quantized_embedding(child, quantization))) + replacements.append((local_path, _float_embedding(child))) continue if quantization is None: continue - if excluded: + if quantization.quant_method not in _AFFINE_METHODS: continue - if is_lm_head and not quantization.quantize_lm_head: - continue - if name.split(".")[-1] in {"router", "shared_expert_gate"}: - continue - if isinstance(child, Embedding) and name.rsplit(".", 1)[-1] in ( - _TOKEN_EMBEDDING_NAMES - ): - if quantization.quantize_embeddings: - if int(child.weight.shape[1]) % quantization.group_size != 0: - raise ValueError( - f"Embedding {name!r} dimension {int(child.weight.shape[1])} " - f"is not divisible by quantization group size " - f"{quantization.group_size}." + + 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, + ), ) - replacements.append((name, _quantized_embedding(child, quantization))) + ) continue + if isinstance(child, Linear) and type(child).forward is Linear.forward: - assert linear_factory is not None out_features, in_features = (int(dim) for dim in child.weight.shape) replacements.append( ( - name, - linear_factory( + local_path, + _linear_factory(config, quantization)( in_features, out_features, bias=child.bias is not None, @@ -486,12 +260,11 @@ def _configure_component_module( ) ) elif type(child) is ClippableLinear: - assert clippable_factory is not None out_features, in_features = (int(dim) for dim in child.weight.shape) replacements.append( ( - name, - clippable_factory( + local_path, + _clippable_linear_factory(config, quantization)( in_features, out_features, bias=child.bias is not None, @@ -499,236 +272,336 @@ def _configure_component_module( ) ) + # 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_module(component_module, path, replacement) + _replace_child(component_module, path, replacement) -def configure_component_quantization( +def _default_manifest( module: nn.Module, config: BaseModelConfig, task: str | ModelTask, -) -> None: - """Configure every task component from ``config.component_quantization``.""" - component_quantization = getattr(config, "component_quantization", None) - if component_quantization is None: - return +) -> ComponentManifest: resolved_task = get_task(task) - paths = _component_module_paths(module, resolved_task) - unresolved = set(component_quantization) - set(paths) - if "model" in paths: + 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 paths: + if "decoder" in manifest: unresolved.discard("model") - if set(paths) == {"model"}: - unresolved -= _KNOWN_COMPONENT_NAMES + if manifest.names == ("model",): + unresolved -= _KNOWN_SPLIT_COMPONENTS if unresolved: raise ValueError( - f"{type(resolved_task).__name__} cannot resolve component module(s) " - f"{sorted(unresolved)} on {type(module).__name__}. Resolved components: " - f"{sorted(paths)}" + f"Component quantization references components not produced by " + f"{type(get_task(task)).__name__}: {sorted(unresolved)}. " + f"Available components: {sorted(manifest)}" ) - for component, path in paths.items(): - quantization = _effective_component_quantization(module, config, component) - if quantization is not None and quantization.quant_method not in _AFFINE_QUANT_METHODS: - component_module = _component_module(module, path) - if not any( - isinstance(child, QuantizedLinear) - for name, child in component_module.named_modules() - if name - ): - raise NotImplementedError( - f"Generic component quantization cannot construct " - f"{quantization.quant_method!r} projections for component " - f"{component!r}; the model must provide a specialized component." - ) + 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(module, path), + component_module, + descriptor, config, quantization, - output_head_paths=_component_output_head_paths(module, component), + owned_by_other_components=owned_elsewhere, ) + return manifest -def _has_raw_packed_weight( - names: Iterable[str], - parameter_names: frozenset[str], -) -> bool: - return any( - name.endswith(("_qweight", ".qweight")) and name not in parameter_names - for name in names - ) +def _raw_qweight_key(name: str) -> bool: + return name.endswith(("_qweight", ".qweight")) -def preprocess_component_quantized_state_dict( - state_dict: dict[str, Any], +def _canonical_component_parameter_keys( module: nn.Module, - config: BaseModelConfig, - task: str | ModelTask, - component_names: Iterable[str], -) -> dict[str, Any]: - """Convert remaining raw packed sidecars with each component's layout.""" - if getattr(config, "component_quantization", None) is None: - return state_dict + descriptor: ComponentDescriptor, +) -> frozenset[str]: + component_module = _resolve_module( + module, + descriptor.module_attribute_path, + ) + if component_module is None: + return frozenset() - resolved_task = get_task(task) - component_paths = _component_module_paths(module, resolved_task) - component_names = tuple(component_names) - parameter_names = frozenset(name for name, _ in module.named_parameters()) - routing_prefixes = { - component: {prefix for prefix in (component, component_paths.get(component)) if prefix} - for component in component_names + 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 routed_component(key: str) -> str | None: + def owner(key: str) -> str | None: matches = [ (len(prefix), component) - for component, prefixes in routing_prefixes.items() - for prefix in prefixes + for component, component_prefixes in prefixes.items() + for prefix in component_prefixes if key.startswith(f"{prefix}.") ] if not matches: - return None + 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." + 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: - component_path = component_paths.get(component, component) - if len(component_names) == 1 or not component_path: - component_weights = dict(result) - else: - component_weights = { - key: value - for key, value in result.items() - if routed_component(key) == component - } - if not component_weights or not _has_raw_packed_weight( - component_weights, - parameter_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 - quantization = _effective_component_quantization(module, config, component) - if quantization is None: - packed_key = next( - key for key in component_weights if key.endswith(("_qweight", ".qweight")) - ) + 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 configured as floating point, but " - f"packed checkpoint weight {packed_key!r} was found." + f"Component {component!r} is floating point but checkpoint " + f"contains packed weight {packed_key!r}" ) - if quantization.quant_method not in _AFFINE_QUANT_METHODS: - raise NotImplementedError( - f"Generic packed-weight preprocessing does not support " - f"{quantization.quant_method!r} for component {component!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}" ) - packed_expert_key = next( - ( - key - for key in component_weights - if key.endswith(("_qweight", ".qweight")) and "expert" in key - ), - None, + + codec = codec_registry.get(component_quantization.quant_method) + bundle = codec.group( + descriptor, + source_weights, + component_quantization, ) - if packed_expert_key is not None: - raise NotImplementedError( - f"Component {component!r} still contains packed expert weight " - f"{packed_expert_key!r} after model preprocessing. Its model " - "must provide a component-aware QMoE conversion." - ) - output_head_paths = _component_output_head_paths(module, component) - packed_lm_head = next( - ( - key - for key in component_weights - if key.endswith(("_qweight", ".qweight")) - and any( - key.startswith(f"{head_path}.") or f".{head_path}." in key - for head_path in output_head_paths + 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." ) - ), - None, - ) - if packed_lm_head is not None and not quantization.quantize_lm_head: - raise ValueError( - f"Component {component!r} keeps lm_head floating point, but " - f"packed checkpoint weight {packed_lm_head!r} was found." - ) - embedding_keys = _component_token_embedding_keys( - module, - component, - component_path, - ) - packed_embeddings = [ - (key, embedding_key) - for key in component_weights - for embedding_key in embedding_keys - if _packed_qweight_for(key, embedding_key) - ] - packed_embedding = packed_embeddings[0][0] if packed_embeddings else None - if packed_embedding is not None and not quantization.quantize_embeddings: - raise ValueError( - f"Component {component!r} keeps embeddings floating point, but " - f"packed checkpoint weight {packed_embedding!r} was found." - ) - if len(packed_embeddings) > 1: - raise NotImplementedError( - f"Component {component!r} contains multiple packed token tables; " - "generic component preprocessing currently supports one per component." - ) - if ( - packed_embedding is not None - and quantization.quant_method != "olive" - and quantization.quantize_embeddings - ): - raise NotImplementedError( - "Generic component preprocessing supports packed token embeddings " - "only for Olive checkpoints." + 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, ) - embed_key = ( - packed_embeddings[0][1] - if packed_embeddings - else embedding_keys[0] - if len(embedding_keys) == 1 - else "__mobius_no_token_embedding__.weight" - ) + 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)) - converted = preprocess_quantized_weights( - component_weights, - quantization, - tie_embeddings=False, - embed_key=embed_key, - qmoe_target_path=None, + canonical_keys = frozenset( + key + for component in component_names + for key in _canonical_component_parameter_keys( + module, + manifest[component], ) - if len(component_names) == 1: - result = converted - else: - for key in component_weights: - result.pop(key, None) - result.update(converted) - remaining_packed_key = next( - ( - key - for key in result - if key.endswith(("_qweight", ".qweight")) and key not in parameter_names - ), + ) + remaining = next( + (key for key in result if _raw_qweight_key(key) and key not in canonical_keys), None, ) - if remaining_packed_key is not None: + if remaining is not None: raise ValueError( - f"Packed checkpoint weight {remaining_packed_key!r} was not routed " - "to any ModelPackage component." + f"Packed checkpoint weight {remaining!r} was not routed to any " + "ModelPackage component" ) return result + + +def attach_hf_component_sources( + module: nn.Module, + *, + model_type: str, + hf_config: object, +) -> None: + """Attach the runtime HF component map selected for this concrete model.""" + resolver = getattr(type(module), "get_hf_component_sources", None) + if resolver is not None: + source_map = resolver(model_type=model_type, hf_config=hf_config) + else: + source_map = getattr(type(module), "HF_COMPONENT_SOURCES", {}) + module._hf_component_sources = { + component: tuple(paths) for component, paths in source_map.items() + } + + +def preprocess_component_quantized_state_dict( + state_dict: dict[str, torch.Tensor], + module: nn.Module, + config: BaseModelConfig, + task: ModelTask | str | None, + package_components: Iterable[str], +) -> dict[str, torch.Tensor]: + """Compatibility wrapper for normalize_component_quantized_weights.""" + import torch + resolved_task = get_task(task) if task is not None else None + return normalize_component_quantized_weights( + state_dict, + module, + config, + package_components, + task=resolved_task, + ) diff --git a/src/mobius/_component_quantization_test.py b/src/mobius/_component_quantization_test.py index 8ee437230..09d2b9f90 100644 --- a/src/mobius/_component_quantization_test.py +++ b/src/mobius/_component_quantization_test.py @@ -1,7 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. -"""Tests for generic per-component quantization wiring.""" +"""Tests for authoritative component quantization plans.""" from __future__ import annotations @@ -9,21 +9,15 @@ import pytest import torch -from onnxscript import OpBuilder, nn +from onnxscript import nn from mobius._component_quantization import ( - attach_hf_component_sources, configure_component_quantization, - preprocess_component_quantized_state_dict, -) -from mobius._configs import ( - ArchitectureConfig, - QuantizationConfig, - QuantizationOverride, + normalize_component_quantized_weights, ) +from mobius._configs import ArchitectureConfig, QuantizationConfig from mobius._model_package import ModelPackage from mobius.components import ( - Embedding, Linear, QuantizedEmbedding, QuantizedLinear, @@ -32,42 +26,55 @@ 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 = Linear): + def __init__(self, linear_class: type[nn.Module] = Linear): super().__init__() self.proj = linear_class(64, 32, bias=False) - def forward(self, op: OpBuilder, x): - return self.proj(op, x) - class _Composite(nn.Module): HF_COMPONENT_SOURCES: ClassVar[dict[str, tuple[str, ...]]] = { - "decoder": ("model.layers",), - "vision_encoder": ("model.visual",), - "audio_encoder": ("model.audio",), - "embedding": ("model.embed_tokens",), + "decoder": ("model.language_model.layers", "lm_head"), + "audio_encoder": ("model.audio_tower",), + "embedding": ("model.language_model.embed_tokens",), } def __init__(self): super().__init__() - quantized = make_quantized_linear_factory(bits=4, block_size=16) - self.decoder = _Projection(quantized) - self.vision_encoder = _Projection() + root_quantized = make_quantized_linear_factory(bits=4, block_size=16) + self.decoder = _Decoder(root_quantized) self.audio_tower = _Projection() - self.embedding = _Projection(quantized) + self.embedding = _Projection(root_quantized) class _CompositeTask(ModelTask): model_roles: ClassVar[dict[str, str]] = { "decoder": "decoder", - "vision_encoder": "encoder", "audio_encoder": "encoder", "embedding": "embedding", } components = ComponentSpec( decoder="decoder", - vision_encoder="vision_encoder", audio_encoder="audio_tower", embedding="embedding", ) @@ -76,250 +83,155 @@ def build(self, module, config) -> ModelPackage: raise NotImplementedError -class _SingleTask(ModelTask): - model_roles: ClassVar[dict[str, str]] = {"model": "decoder"} - - 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, - "vision_encoder": QuantizationConfig( + "audio_encoder": QuantizationConfig( bits=8, group_size=32, quant_method="olive", sym=True, ), - "audio_encoder": QuantizationConfig( - bits=2, - group_size=16, - quant_method="olive", - sym=True, - ), }, ) -def test_configures_quantized_and_float_components_independently(): +def test_component_plan_applies_regex_exclusions_per_linear(): module = _Composite() configure_component_quantization(module, _config(), _CompositeTask()) - assert isinstance(module.decoder.proj, QuantizedLinear) - assert (module.decoder.proj._bits, module.decoder.proj._block_size) == (4, 16) - assert isinstance(module.vision_encoder.proj, QuantizedLinear) - assert (module.vision_encoder.proj._bits, module.vision_encoder.proj._block_size) == ( - 8, - 32, - ) + 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) == ( - 2, - 16, - ) - assert type(module.embedding.proj) is Linear - - -def test_dynamic_hf_sources_drive_component_override_layout(): - class _DynamicComposite(_Composite): - @classmethod - def get_hf_component_sources( - cls, - *, - model_type: str, - hf_config: object, - ) -> dict[str, tuple[str, ...]]: - assert model_type == "alternate" - assert hf_config is not None - return { - **cls.HF_COMPONENT_SOURCES, - "vision_encoder": ("model.vision_model", "model.connector"), - } - - module = _DynamicComposite() - attach_hf_component_sources( - module, - model_type="alternate", - hf_config=object(), - ) - quantization = QuantizationConfig( - bits=4, - group_size=16, - quant_method="olive", - overrides={ - "model.vision_model": QuantizationOverride(bits=8, group_size=32), - "model.connector": QuantizationOverride(bits=8, group_size=32), - }, - ) - config = ArchitectureConfig( - quantization=quantization, - component_quantization={"vision_encoder": quantization}, - ) - - configure_component_quantization(module, config, _CompositeTask()) - - assert isinstance(module.vision_encoder.proj, QuantizedLinear) - assert (module.vision_encoder.proj._bits, module.vision_encoder.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_declared_output_head_respects_quantize_lm_head_flag(): - class _Decoder(nn.Module): - def __init__(self): - super().__init__() - self.proj = Linear(64, 32, bias=False) - self.proj_out = Linear(64, 256, bias=False) - - class _Model(nn.Module): - HF_COMPONENT_SOURCES: ClassVar[dict[str, tuple[str, ...]]] = { - "decoder": ("model.decoder", "proj_out") - } - COMPONENT_OUTPUT_HEADS: ClassVar[dict[str, tuple[str, ...]]] = { - "decoder": ("proj_out",) - } - - def __init__(self): - super().__init__() - self.decoder = _Decoder() - - class _Task(ModelTask): - model_roles: ClassVar[dict[str, str]] = {"decoder": "decoder"} - components: ClassVar[ComponentSpec] = ComponentSpec(decoder="decoder") - - def build(self, module, config) -> ModelPackage: - raise NotImplementedError +def test_specialized_quantized_subclass_fails_instead_of_losing_semantics(): + class _SpecialQuantizedLinear(QuantizedLinear): + def forward(self, op, x): + return super().forward(op, x) - quantization = QuantizationConfig( - bits=4, - group_size=16, - quant_method="olive", - quantize_lm_head=False, - ) + module = _Projection(_SpecialQuantizedLinear) config = ArchitectureConfig( - quantization=quantization, - component_quantization={"decoder": quantization}, + component_quantization={ + "model": QuantizationConfig( + bits=8, + group_size=32, + quant_method="olive", + ) + } ) - module = _Model() - - configure_component_quantization(module, config, _Task()) - - assert isinstance(module.decoder.proj, QuantizedLinear) - assert type(module.decoder.proj_out) is Linear - with pytest.raises(ValueError, match="keeps lm_head floating point"): - preprocess_component_quantized_state_dict( - { - "decoder.proj_out.weight_qweight": torch.zeros(256, 32, dtype=torch.uint8), - "decoder.proj_out.weight_scales": torch.ones(256, 4), - }, - module, - config, - _Task(), - ("decoder",), - ) + + with pytest.raises(TypeError, match="specialized quantized module"): + configure_component_quantization(module, config, _SingleTask()) -def test_preprocesses_raw_weights_with_component_layouts(): +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.proj.weight_qweight": torch.zeros(32, 32, dtype=torch.uint8), - "decoder.proj.weight_scales": torch.ones(32, 4), - "vision_encoder.proj.weight_qweight": torch.zeros(32, 64, dtype=torch.uint8), - "vision_encoder.proj.weight_scales": torch.ones(32, 2), - "audio_tower.proj.weight_qweight": torch.zeros(32, 16, dtype=torch.uint8), - "audio_tower.proj.weight_scales": torch.ones(32, 4), + "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 = preprocess_component_quantized_state_dict( + result = normalize_component_quantized_weights( state_dict, module, config, - _CompositeTask(), - ("decoder", "vision_encoder", "audio_encoder", "embedding"), + ("decoder", "audio_encoder", "embedding"), + manifest=manifest, + task=task, ) - assert result["decoder.proj.weight"].shape == (32, 4, 8) - assert result["vision_encoder.proj.weight"].shape == (32, 2, 32) - assert result["audio_tower.proj.weight"].shape == (32, 4, 4) - assert result["embedding.proj.weight"].shape == (32, 64) - - -def test_single_graph_uses_decoder_layout_and_ignores_split_metadata(): - decoder = QuantizationConfig( - bits=4, - group_size=16, - quant_method="olive", - sym=True, + 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, ) - config = ArchitectureConfig( - quantization=decoder, - component_quantization={ - "decoder": decoder, - "vision_encoder": QuantizationConfig( - bits=8, - group_size=32, - quant_method="olive", - ), - }, - ) - module = _Projection() - configure_component_quantization(module, config, _SingleTask()) - assert isinstance(module.proj, QuantizedLinear) - assert (module.proj._bits, module.proj._block_size) == (4, 16) +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, + ) -def test_quantize_embeddings_only_rewrites_input_token_table(): - class _EmbeddingModule(nn.Module): - def __init__(self): - super().__init__() - self.embed_tokens = Embedding(256, 64) - self.embed_positions = Embedding(128, 64) - quantization = QuantizationConfig( - bits=4, - group_size=16, - quant_method="olive", - quantize_embeddings=True, - ) - config = ArchitectureConfig( - quantization=quantization, - component_quantization={"model": quantization}, - ) - module = _EmbeddingModule() +class _SingleTask(ModelTask): + model_roles: ClassVar[dict[str, str]] = {"model": "decoder"} - configure_component_quantization(module, config, _SingleTask()) + def build(self, module, config) -> ModelPackage: + raise NotImplementedError - assert isinstance(module.embed_tokens, QuantizedEmbedding) - assert type(module.embed_positions) is Embedding +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_existing_quantized_embedding_retargets_component_layout(): - class _EmbeddingModule(nn.Module): - def __init__(self): - super().__init__() - self.embed_tokens = QuantizedEmbedding( - 256, - 64, - bits=8, - block_size=32, - has_zero_point=True, - ) +def test_canonical_quantized_embedding_is_not_treated_as_raw_sidecars(): + module = _QuantizedEmbeddingModel() quantization = QuantizationConfig( bits=4, group_size=16, @@ -331,148 +243,26 @@ def __init__(self): quantization=quantization, component_quantization={"model": quantization}, ) - module = _EmbeddingModule() - original = module.embed_tokens - - configure_component_quantization(module, config, _SingleTask()) - - assert module.embed_tokens is not original - assert isinstance(module.embed_tokens, QuantizedEmbedding) - assert (module.embed_tokens._bits, module.embed_tokens._block_size) == (4, 16) - assert module.embed_tokens.zero_points is None - - -def test_projection_name_containing_embedding_is_not_a_token_table(): - class _Embeddings(nn.Module): - def __init__(self): - super().__init__() - self.patch_embedding = Linear(64, 32, bias=False) - - class _VisionProjection(nn.Module): - def __init__(self): - super().__init__() - self.embeddings = _Embeddings() - - quantization = QuantizationConfig( - bits=4, - group_size=16, - quant_method="olive", - quantize_embeddings=False, - ) - config = ArchitectureConfig( - quantization=quantization, - component_quantization={"model": quantization}, - ) - module = _VisionProjection() - configure_component_quantization(module, config, _SingleTask()) - - result = preprocess_component_quantized_state_dict( - { - "embeddings.patch_embedding.weight_qweight": torch.zeros( - 32, 32, dtype=torch.uint8 - ), - "embeddings.patch_embedding.weight_scales": torch.ones(32, 4), - }, - module, - config, - _SingleTask(), - ("model",), - ) - - assert result["embeddings.patch_embedding.weight"].shape == (32, 4, 8) - - -def test_nonstandard_token_embedding_name_keeps_olive_table_2d(): - class _WordEmbeddingModule(nn.Module): - def __init__(self): - super().__init__() - self.word_embeddings = Embedding(256, 64) + 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), + } - quantization = QuantizationConfig( - bits=4, - group_size=16, - quant_method="olive", - quantize_embeddings=True, - ) - config = ArchitectureConfig( - quantization=quantization, - component_quantization={"model": quantization}, - ) - module = _WordEmbeddingModule() - configure_component_quantization(module, config, _SingleTask()) - - qweight = torch.zeros(256, 32, dtype=torch.uint8) - result = preprocess_component_quantized_state_dict( - { - "word_embeddings.weight_qweight": qweight, - "word_embeddings.weight_scales": torch.ones(256, 4), - }, + result = normalize_component_quantized_weights( + state_dict, module, config, - _SingleTask(), ("model",), + manifest=manifest, + task=task, ) - assert result["word_embeddings.qweight"] is qweight - assert result["word_embeddings.qweight"].ndim == 2 - - -def test_scaled_embedding_quantization_preserves_forward_semantics(): - class _ScaledEmbedding(Embedding): - def __init__(self): - super().__init__(256, 64) - self.embed_scale = 2.0 - - def forward(self, op, input_ids): - return op.Mul(super().forward(op, input_ids), self.embed_scale) - - class _ScaledEmbeddingModule(nn.Module): - def __init__(self): - super().__init__() - self.embed_tokens = _ScaledEmbedding() - - quantization = QuantizationConfig( - bits=4, - group_size=16, - quant_method="olive", - quantize_embeddings=True, - ) - config = ArchitectureConfig( - quantization=quantization, - component_quantization={"model": quantization}, - ) - - module = _ScaledEmbeddingModule() - configure_component_quantization(module, config, _SingleTask()) - - assert isinstance(module.embed_tokens, QuantizedEmbedding) - assert module.embed_tokens.embed_scale == pytest.approx(2.0) - - -def test_unknown_specialized_embedding_fails_before_losing_forward_semantics(): - class _SpecialEmbedding(Embedding): - def forward(self, op, input_ids): - return op.Neg(super().forward(op, input_ids)) - - class _SpecialEmbeddingModule(nn.Module): - def __init__(self): - super().__init__() - self.embed_tokens = _SpecialEmbedding(256, 64) - - quantization = QuantizationConfig( - bits=4, - group_size=16, - quant_method="olive", - quantize_embeddings=True, - ) - config = ArchitectureConfig( - quantization=quantization, - component_quantization={"model": quantization}, - ) - - with pytest.raises(TypeError, match="specialized embedding"): - configure_component_quantization( - _SpecialEmbeddingModule(), - config, - _SingleTask(), - ) + 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/_base.py b/src/mobius/_configs/_base.py index f2ad90233..9a89c145b 100644 --- a/src/mobius/_configs/_base.py +++ b/src/mobius/_configs/_base.py @@ -390,7 +390,6 @@ def _parse_component_quantization_mapping( *, expert_dtype: object | None, ) -> dict[str, QuantizationConfig]: - """Parse an explicit component-name to quantization-config mapping.""" if not isinstance(value, Mapping): raise TypeError( "component_quantization must be a mapping of component names to " @@ -421,14 +420,12 @@ def _extract_component_quantization( parent_config: object | None, decoder_quantization: QuantizationConfig | None, ) -> dict[str, QuantizationConfig] | None: - """Extract explicit or nested per-component quantization metadata.""" + """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) - # Explicit mapping is authoritative. Accept a top-level field or a - # ``components`` mapping nested in the traditional quantization_config. for source in sources: declaration = _get_config_value(source, "component_quantization") if declaration is None: @@ -445,8 +442,6 @@ def _extract_component_quantization( expert_dtype=_get_config_value(source, "expert_dtype"), ) - # Composite checkpoints may instead put independent quantization_config - # dictionaries directly on their vision/audio sub-configs. composite = parent_config or config nested: dict[str, QuantizationConfig] = {} found_nested_declaration = False @@ -500,8 +495,8 @@ class BaseModelConfig: # Model dtype (from HF config dtype) dtype: ir.DataType = ir.DataType.FLOAT quantization: QuantizationConfig | None = None - # ``None`` preserves the legacy model-wide quantization behavior. A mapping - # is authoritative: omitted components remain floating point. + # ``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. @@ -512,7 +507,7 @@ class BaseModelConfig: diffusion_shift_logits: bool = False def quantization_for(self, component: str) -> QuantizationConfig | None: - """Return the effective quantization config for one package component.""" + """Return the effective quantization plan for one package component.""" if self.component_quantization is None: return self.quantization candidates = { diff --git a/src/mobius/_configs/_quantization.py b/src/mobius/_configs/_quantization.py index a44e3a576..0d8b8166e 100644 --- a/src/mobius/_configs/_quantization.py +++ b/src/mobius/_configs/_quantization.py @@ -6,12 +6,20 @@ 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 overrides emitted by Olive mixed-precision quantization.""" + """Per-module affine layout override emitted by an upstream quantizer.""" bits: int | None = None group_size: int | None = None @@ -19,7 +27,7 @@ class QuantizationOverride: @classmethod def from_value(cls, value: object) -> QuantizationOverride: - """Parse one serialized Olive override.""" + """Parse one serialized module override.""" if isinstance(value, cls): return value if hasattr(value, "to_dict"): @@ -35,7 +43,7 @@ def from_value(cls, value: object) -> QuantizationOverride: ) def apply(self, config: QuantizationConfig) -> QuantizationConfig: - """Return *config* with this override applied.""" + """Return *config* with this module override applied.""" updates = { name: value for name, value in dataclasses.asdict(self).items() @@ -76,13 +84,11 @@ 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 - # Root-relative HuggingFace module names left in floating point by Olive's - # mixed-precision planner. ``None`` means no component plan was recorded; - # an empty tuple means the planner explicitly quantized every eligible - # module with the default configuration. - modules_to_not_convert: tuple[str, ...] | None = None - # Olive per-module precision overrides. Mobius can collapse uniform - # overrides beneath one component into that component's effective config. + # 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 @@ -92,7 +98,7 @@ def from_value( *, expert_dtype: object | None = None, ) -> QuantizationConfig | None: - """Parse a serialized HuggingFace quantization-config value.""" + """Parse one serialized HuggingFace quantization configuration.""" if value is None: return None if isinstance(value, cls): @@ -176,16 +182,22 @@ def from_value( # fp8 was already routed to the typed blocker above.) if method == "fp8": return None - raw_modules_to_not_convert = qc.get("modules_to_not_convert") - if raw_modules_to_not_convert is not None and not isinstance( - raw_modules_to_not_convert, (list, tuple) - ): + 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), @@ -196,15 +208,8 @@ def from_value( 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=( - tuple(str(name) for name in raw_modules_to_not_convert) - if raw_modules_to_not_convert is not None - else None - ), - overrides={ - str(name): QuantizationOverride.from_value(override) - for name, override in raw_overrides.items() - }, + modules_to_not_convert=exclusions, + overrides=overrides, ) @classmethod @@ -220,145 +225,48 @@ def from_transformers(cls, hf_config) -> QuantizationConfig | None: @property def has_module_plan(self) -> bool: - """Whether Olive recorded component-selection metadata.""" - return self.modules_to_not_convert is not None or bool(self.overrides) + """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:"): + 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 or module_name.startswith(f"{pattern}.") or pattern in 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 def for_source_paths( self, source_paths: tuple[str, ...], *, - component: str, + component: str = "decoder", ) -> QuantizationConfig | None: - """Collapse a uniform Olive module plan into one component config. - - A component-level ONNX graph cannot represent different quantization - layouts for individual projections without constructing each projection - separately. Uniform overrides are therefore accepted, fully excluded - components stay float, and mixed layouts fail loudly. - """ - if not source_paths: - raise ValueError( - f"Cannot derive component quantization for {component!r}: " - "the model declares no HuggingFace source paths." - ) - - def targets_component(name: str) -> bool: - return any( - name == prefix - or name.startswith(f"{prefix}.") - or prefix.startswith(f"{name}.") - for prefix in source_paths - ) - - def covers_source_path(name: str, source_path: str) -> bool: - return name == source_path or source_path.startswith(f"{name}.") - - regex_exclusions = [ - name for name in self.modules_to_not_convert or () if name.startswith("re:") - ] - regex_overrides = [name for name in self.overrides if name.startswith("re:")] - if regex_exclusions or regex_overrides: - raise ValueError( - f"Cannot derive component quantization for {component!r} from " - "regex module rules. Store an explicit component_quantization " - "mapping in the checkpoint config instead." - ) - - exclusions = [ - name for name in self.modules_to_not_convert or () if targets_component(name) - ] - component_overrides = [ - (name, override) - for name, override in self.overrides.items() - if targets_component(name) - ] - if exclusions and component_overrides: - raise ValueError( - f"Component {component!r} mixes excluded and quantized modules; " - "Mobius requires one quantization configuration per component." - ) - if exclusions: - fully_excluded = all( - any(covers_source_path(name, source_path) for name in exclusions) - for source_path in source_paths - ) - if fully_excluded: - return None - raise ValueError( - f"Component {component!r} has partial module exclusions; " - "Mobius requires one quantization configuration per component." - ) - - base_layout = (self.bits, self.group_size, self.sym) - effective = [ - ( - override.bits if override.bits is not None else self.bits, - (override.group_size if override.group_size is not None else self.group_size), - override.sym if override.sym is not None else self.sym, - ) - for _, override in component_overrides - ] - different_layouts = {layout for layout in effective if layout != base_layout} - if len(different_layouts) > 1: - raise ValueError( - f"Component {component!r} has multiple quantization layouts " - f"{sorted(different_layouts | {base_layout})!r}; " - "Mobius requires one per component." - ) - - config = self - if different_layouts: - target_layout = next(iter(different_layouts)) - fully_overridden = all( - any( - covers_source_path(name, source_path) and effective[index] == target_layout - for index, (name, _) in enumerate(component_overrides) - ) - for source_path in source_paths - ) - if not fully_overridden: - raise ValueError( - f"Component {component!r} mixes the default layout " - f"{base_layout!r} with override layout {target_layout!r}; " - "store an explicit component_quantization mapping instead." - ) - override = next( - override - for index, (_, override) in enumerate(component_overrides) - if effective[index] == target_layout - ) - config = override.apply(self) - return dataclasses.replace( - config, - modules_to_not_convert=None, - overrides={}, - ) - - def for_components( - self, - component_sources: Mapping[str, tuple[str, ...]], - ) -> dict[str, QuantizationConfig]: - """Collapse an Olive module plan into package-component layouts.""" - result: dict[str, QuantizationConfig] = {} - for component, source_paths in component_sources.items(): - if not source_paths: - continue - effective_paths = source_paths - if component in {"decoder", "model"}: - # LM-head and token-table selection already have dedicated - # QuantizationConfig flags. Their exclusions must not turn an - # otherwise quantized decoder backbone into a float component. - effective_paths = tuple( - path - for path in source_paths - if not path.endswith(("lm_head", "embed_tokens")) - ) - if not effective_paths: - effective_paths = source_paths - quantization = self.for_source_paths( - effective_paths, - component=component, - ) - if quantization is not None: - result[component] = quantization - return result + """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/_configs_test.py b/src/mobius/_configs_test.py index 8c1193e54..8678c556e 100644 --- a/src/mobius/_configs_test.py +++ b/src/mobius/_configs_test.py @@ -18,7 +18,6 @@ GlmAsrConfig, MuseGlimmerConfig, QuantizationConfig, - QuantizationOverride, VisionConfig, _extract_audio_config, _extract_mrope_fields, @@ -1059,64 +1058,37 @@ def test_from_transformers_olive_component_flags(self): assert qc.quantize_lm_head is True assert qc.quantize_vision is True - def test_from_transformers_olive_module_plan(self): - hf = SimpleNamespace( - quantization_config={ + def test_component_plan_matches_exact_and_regex_module_rules(self): + qc = QuantizationConfig.from_value( + { "quant_method": "olive", "bits": 4, "group_size": 32, - "symmetric": False, - "modules_to_not_convert": ["model.embed_tokens"], + "modules_to_not_convert": [ + r"re:.*\.per_layer_input_gate", + ], "overrides": { - "model.vision_tower": { + "model.layers.0.q_proj": { "bits": 8, "group_size": 64, - "symmetric": True, } }, } ) - qc = QuantizationConfig.from_transformers(hf) - assert qc is not None - assert qc.modules_to_not_convert == ("model.embed_tokens",) - assert qc.has_module_plan is True - vision = qc.for_source_paths( - ("model.vision_tower",), - component="vision_encoder", - ) - assert vision is not None - assert (vision.bits, vision.group_size, vision.sym) == (8, 64, True) - assert vision.modules_to_not_convert is None - assert vision.overrides == {} - - def test_module_plan_keeps_excluded_component_float(self): - qc = QuantizationConfig( - quant_method="olive", - modules_to_not_convert=("model.audio_tower",), - ) - - assert ( - qc.for_source_paths( - ("model.audio_tower",), - component="audio_encoder", - ) - is None - ) - - def test_module_plan_rejects_partial_component_override(self): - qc = QuantizationConfig( - bits=4, - group_size=32, - quant_method="olive", - overrides={"model.audio_tower.layers.0.q_proj": QuantizationOverride(bits=8)}, - ) - - with pytest.raises(ValueError, match="mixes the default layout"): - qc.for_source_paths( - ("model.audio_tower",), - component="audio_encoder", + 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): @@ -1134,21 +1106,19 @@ def test_architecture_config_parses_explicit_component_quantization(self): parent = SimpleNamespace( model_type="composite", component_quantization={ - "text": { + "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, }, - "audio": { - "quant_method": "gptq", - "bits": 2, - "group_size": 16, - }, }, ) @@ -1157,51 +1127,7 @@ def test_architecture_config_parses_explicit_component_quantization(self): 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("audio_encoder").bits == 2 - assert config.quantization is config.quantization_for("decoder") - - def test_architecture_config_parses_nested_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, - quantization_config={ - "quant_method": "olive", - "bits": 4, - "group_size": 32, - }, - ) - parent = SimpleNamespace( - model_type="composite", - vision_config=SimpleNamespace( - quantization_config={ - "quant_method": "olive", - "bits": 8, - "group_size": 64, - } - ), - audio_config=SimpleNamespace( - quantization_config={ - "quant_method": "olive", - "bits": 2, - "group_size": 16, - } - ), - ) - - 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("embedding").bits == 4 - assert config.quantization_for("vision_encoder").bits == 8 - assert config.quantization_for("audio_encoder").bits == 2 + assert config.quantization_for("decoder").modules_to_not_convert def test_quantize_component_flags_default_false(self): qc = QuantizationConfig() diff --git a/src/mobius/integrations/transformers/_builder.py b/src/mobius/integrations/transformers/_builder.py index d45495a5b..005547f1d 100644 --- a/src/mobius/integrations/transformers/_builder.py +++ b/src/mobius/integrations/transformers/_builder.py @@ -15,7 +15,7 @@ from mobius._builder import build_from_module, resolve_dtype from mobius._component_quantization import ( attach_hf_component_sources, - preprocess_component_quantized_state_dict, + normalize_component_quantized_weights, ) from mobius._model_package import ModelPackage from mobius._registry import registry @@ -507,6 +507,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) attach_hf_component_sources( model_module, @@ -522,6 +530,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}" @@ -587,12 +596,13 @@ 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 = preprocess_component_quantized_state_dict( + state_dict = normalize_component_quantized_weights( state_dict, model_module, config, - task, package.keys(), + manifest=component_manifest, + task=resolved_task, ) package.apply_weights( state_dict, diff --git a/src/mobius/integrations/transformers/_builder_test.py b/src/mobius/integrations/transformers/_builder_test.py index 85a72542d..831fc4e9a 100644 --- a/src/mobius/integrations/transformers/_builder_test.py +++ b/src/mobius/integrations/transformers/_builder_test.py @@ -305,10 +305,15 @@ def test_strip_to_text_only_drops_component_quantization() -> None: quant_method="olive", ) config = make_config( + quantization=decoder, component_quantization={ "decoder": decoder, - "vision_encoder": decoder, - } + "vision_encoder": QuantizationConfig( + bits=8, + group_size=32, + quant_method="olive", + ), + }, ) stripped = transformers_builder._strip_to_text_only(config, "qwen2") @@ -342,8 +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 == {} - - 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 8010e5520..5338d8a62 100644 --- a/src/mobius/models/gemma4_test.py +++ b/src/mobius/models/gemma4_test.py @@ -586,6 +586,39 @@ def test_component_override_uses_same_layout_for_graph_and_weights(self): ) +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/tests/build_graph/_support.py b/tests/build_graph/_support.py index e5c932876..b5156c83f 100644 --- a/tests/build_graph/_support.py +++ b/tests/build_graph/_support.py @@ -63,6 +63,37 @@ def _run_onnx_checker(pkg: dict[str, ir.Model], model_type: str) -> None: _onnx_checker(model) +def _with_component_quantization(config, task): + """Assign distinct tiny affine layouts to every materialized component.""" + import dataclasses + from mobius._configs import QuantizationConfig + + 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 getattr(config, "audio", None) 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, + ) + + def _assert_outputs_have_shapes_and_dtypes( pkg: dict[str, ir.Model], model_type: str, diff --git a/tests/build_graph/vision_language_test.py b/tests/build_graph/vision_language_test.py index 7315b9380..1cbcc596f 100644 --- a/tests/build_graph/vision_language_test.py +++ b/tests/build_graph/vision_language_test.py @@ -25,6 +25,7 @@ _assert_outputs_have_shapes_and_dtypes, _make_params, _run_onnx_checker, + _with_component_quantization, ) from mobius._builder import DTYPE_MAP, build_from_module from mobius._configs import (