Load quantized weights per component

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
 

src/mobius/__main__.py

@@ -521,6 +521,15 @@
             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 @@
             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 @@
                 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:

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):