Migrate composite model weight adapters

CHANGELOG.md

@@ -19,6 +19,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   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.
+- Gemma4, Qwen3.5/QMoE, and T5 adapters preserve their architecture-specific
+  rename, tied-weight, and expert-packing semantics while handing independently
+  routed component sidecars to the generic codec and binding validator.
 
 ### Packed fused MoE experts (Olive/GPTQ/AWQ) survive HF weight renaming
 

src/mobius/__main__.py

@@ -452,6 +452,10 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask:
         for name, model in pkg.items():
             model.graph.name = f"{config_path}/{name}"
         if load_weights:
+            from mobius._component_quantization import (
+                validate_quantized_component_bindings,
+            )
+
             if compressed_tensors_config is not None:
                 # Packed FP4 weights cannot pass through ordinary apply_weights.
                 # The same loader owns both faithful native storage and the
@@ -475,8 +479,14 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask:
                 )
             else:
                 state_dict = _load_weights_from_dir(config_path)
-                if hasattr(model_module, "preprocess_weights"):
-                    state_dict = model_module.preprocess_weights(state_dict)
+                from mobius.weights import adapt_model_weights
+
+                state_dict = adapt_model_weights(
+                    model_module,
+                    state_dict,
+                    config=config,
+                    manifest=component_manifest,
+                )
                 from mobius._component_quantization import (
                     normalize_component_quantized_weights,
                 )
@@ -490,6 +500,7 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask:
                     task=resolved_task,
                 )
                 pkg.apply_weights(state_dict)
+            validate_quantized_component_bindings(pkg, config)
     else:
         model_id_or_path = args.model
         if static_cache_params is not None:

src/mobius/_component_quantization.py

@@ -8,6 +8,7 @@
 __all__ = [
     "configure_component_quantization",
     "normalize_component_quantized_weights",
+    "validate_quantized_component_bindings",
 ]
 
 from collections.abc import Iterable, Mapping
@@ -127,17 +128,63 @@ def _float_embedding(module: QuantizedEmbedding) -> Embedding:
     )
 
 
+def _linear_layout_matches(
+    module: QuantizedLinear,
+    quantization: QuantizationConfig,
+) -> bool:
+    expected_zero_point_dtype = (
+        module.scales.dtype if quantization.float_zero_point else ir.DataType.UINT8
+    )
+    return (
+        module._bits == quantization.bits
+        and module._block_size == quantization.group_size
+        and (module.zero_points is None) is quantization.sym
+        and (
+            module.zero_points is None or module.zero_points.dtype == expected_zero_point_dtype
+        )
+    )
+
+
+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 _effective_module_quantization(
     component_quantization: QuantizationConfig | None,
     descriptor: ComponentDescriptor,
     local_module_path: str,
+    *,
+    source_module_names: tuple[str, ...] | None = None,
 ) -> QuantizationConfig | None:
     if component_quantization is None or component_quantization.quant_method == "none":
         return None
-    source_names = descriptor.source_module_names(local_module_path)
+    source_names = (
+        source_module_names
+        if source_module_names is not None
+        else descriptor.source_module_names(local_module_path)
+    )
     return component_quantization.for_module(source_names)
 
 
+def _source_module_names(
+    descriptor: ComponentDescriptor,
+    local_module_path: str,
+    module: nn.Module,
+) -> tuple[str, ...]:
+    names = descriptor.source_module_names(local_module_path)
+    if isinstance(module, (ClippableLinear, ClippableQuantizedLinear)):
+        names = (*names, *(f"{name}.linear" for name in names))
+    return tuple(dict.fromkeys(names))
+
+
 def _configure_component_module(
     component_module: nn.Module,
     descriptor: ComponentDescriptor,
@@ -159,6 +206,11 @@ def _configure_component_module(
             component_quantization,
             descriptor,
             local_path,
+            source_module_names=_source_module_names(
+                descriptor,
+                local_path,
+                child,
+            ),
         )
         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:
@@ -190,6 +242,11 @@ def _configure_component_module(
 
         if isinstance(child, QuantizedLinear):
             if type(child).forward is not QuantizedLinear.forward:
+                if quantization is not None and _linear_layout_matches(
+                    child,
+                    quantization,
+                ):
+                    continue
                 raise TypeError(
                     f"Component plan cannot rewrite specialized quantized "
                     f"module {local_path!r} ({type(child).__name__}); provide "
@@ -209,6 +266,11 @@ def _configure_component_module(
 
         if isinstance(child, QuantizedEmbedding):
             if type(child).forward is not QuantizedEmbedding.forward:
+                if quantization is not None and _embedding_layout_matches(
+                    child,
+                    quantization,
+                ):
+                    continue
                 raise TypeError(
                     f"Component plan cannot rewrite specialized quantized "
                     f"embedding {local_path!r} ({type(child).__name__}); "
@@ -537,10 +599,25 @@ def normalize_component_quantized_weights(
                     "component-specific tied-weight adapter."
                 )
             local_path = _local_weight_module_path(record.name, descriptor)
+            component_module = _resolve_module(
+                module,
+                descriptor.module_attribute_path,
+            )
+            local_module = (
+                _resolve_module(component_module, local_path)
+                if component_module is not None
+                else None
+            )
+            source_names = (
+                _source_module_names(descriptor, local_path, local_module)
+                if local_module is not None
+                else descriptor.source_module_names(local_path)
+            )
             quantization = _effective_module_quantization(
                 component_quantization,
                 descriptor,
                 local_path,
+                source_module_names=source_names,
             )
             if quantization is None:
                 raise ValueError(
@@ -567,3 +644,35 @@ def normalize_component_quantized_weights(
             "ModelPackage component"
         )
     return result
+
+
+def validate_quantized_component_bindings(
+    models: Mapping[str, ir.Model],
+    config: BaseModelConfig,
+) -> None:
+    """Require every affine quantized op input to carry a bound value."""
+    if getattr(config, "component_quantization", None) is None:
+        return
+
+    quantized_input_slots = {
+        "MatMulNBits": (1, 2, 3),
+        "GatherBlockQuantized": (0, 2, 3),
+    }
+    for component, model in models.items():
+        if _component_quantization(config, component) is None:
+            continue
+        for node in ir.traversal.RecursiveGraphIterator(model.graph):
+            slots = quantized_input_slots.get(node.op_type)
+            if slots is None:
+                continue
+            for index in slots:
+                if index >= len(node.inputs):
+                    continue
+                value = node.inputs[index]
+                if value is None or value.producer() is not None:
+                    continue
+                if value.const_value is None:
+                    raise ValueError(
+                        f"Quantized component {component!r} has unbound "
+                        f"{node.op_type} parameter {value.name!r}"
+                    )

src/mobius/_component_quantization_test.py

@@ -14,6 +14,7 @@
 from mobius._component_quantization import (
     configure_component_quantization,
     normalize_component_quantized_weights,
+    validate_quantized_component_bindings,
 )
 from mobius._configs import ArchitectureConfig, QuantizationConfig
 from mobius._model_package import ModelPackage
@@ -266,3 +267,34 @@ def test_canonical_quantized_embedding_is_not_treated_as_raw_sidecars():
     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"]
+
+
+def test_binding_validator_rejects_unfilled_quantized_parameters():
+    from mobius._testing import create_test_builder, create_test_input
+    from mobius.tasks._base import _make_model
+
+    linear = QuantizedLinear(
+        64,
+        32,
+        bits=4,
+        block_size=16,
+        has_zero_point=False,
+    )
+    builder, op, graph = create_test_builder()
+    x = create_test_input(builder, "x", [1, 64])
+    output = linear(op, x)
+    builder._adapt_outputs([output], "")
+    quantization = QuantizationConfig(
+        bits=4,
+        group_size=16,
+        quant_method="olive",
+    )
+
+    with pytest.raises(ValueError, match="unbound MatMulNBits parameter"):
+        validate_quantized_component_bindings(
+            {"model": _make_model(graph)},
+            ArchitectureConfig(
+                quantization=quantization,
+                component_quantization={"model": quantization},
+            ),
+        )

src/mobius/components/_moe.py

@@ -429,10 +429,15 @@ def __init__(
         assert config.num_experts_per_tok is not None
         self.num_experts = config.num_local_experts
         self.top_k = config.num_experts_per_tok
+        quantization = (
+            config.quantization_for("decoder")
+            if config.component_quantization is not None
+            else config.quantization
+        )
         self._qmoe_quantization = (
             None
             if getattr(config, "disable_qmoe", False)
-            else _supported_qmoe_quantization(config.quantization)
+            else _supported_qmoe_quantization(quantization)
         )
         # Clipped-SwiGLU attributes (QMoE's ``activation_alpha``/``activation_beta``/
         # ``swiglu_limit``). Left ``None`` by default so existing callers get a

src/mobius/integrations/transformers/_builder.py

@@ -13,7 +13,10 @@
 from onnxscript import nn
 
 from mobius._builder import build_from_module, resolve_dtype
-from mobius._component_quantization import normalize_component_quantized_weights
+from mobius._component_quantization import (
+    normalize_component_quantized_weights,
+    validate_quantized_component_bindings,
+)
 from mobius._model_package import ModelPackage
 from mobius._registry import registry
 from mobius.integrations._weight_loading import (
@@ -26,6 +29,7 @@
     stream_compressed_tensors_to_package,
 )
 from mobius.tasks import ModelTask
+from mobius.weights import adapt_model_weights
 
 logger = logging.getLogger(__name__)
 
@@ -410,8 +414,12 @@ def build_transformers_model(
             )
         else:
             state_dict = _download_weights(model_id, revision=revision)
-            if hasattr(model_module, "preprocess_weights"):
-                state_dict = model_module.preprocess_weights(state_dict)
+            state_dict = adapt_model_weights(
+                model_module,
+                state_dict,
+                config=config,
+                manifest=component_manifest,
+            )
             state_dict = normalize_component_quantized_weights(
                 state_dict,
                 model_module,
@@ -424,6 +432,7 @@ def build_transformers_model(
                 state_dict,
                 prefix_map=getattr(model_module, "weight_prefix_map", None),
             )
+        validate_quantized_component_bindings(package, config)
     return package