Validate QMoE scales and defer FC1 interleaving

Reject invalid legacy uniform INT4 scales and serialize FC1 transforms without caching duplicate checkpoint payloads. Preserve component-local transform routing and cover asymmetric zero points and distinct multilayer payloads.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot-Session: fba8a408-cd63-49f0-b8b4-5daa19ed42d4
Signed-off-by: titaiwang <titaiwang@microsoft.com>

Author: titaiwangms

README.md

@@ -78,11 +78,20 @@ FC scale dtypes must match the activations. Every projection's scales must be
 finite and nonnegative. When the adapter supplies the activation dtype,
 scales must remain finite after conversion. Non-INT3 layouts retain the
 standard weight loader's dtype conversion.
+The legacy uniform INT4 path receives the same numeric scale checks.
 INT2, INT4, and INT8 projections retain their existing zero-point
 semantics; asymmetric zero points must be packed `uint8` sidecars.
 Dense INT3 `MatMulNBits`, embeddings, other source layouts, and schema-unapproved
 tails are not enabled. Keep non-expert projections at supported precisions.
 
+`build()` defers INT3-layout FC1 interleaving to initializer serialization,
+using uncached lazy tensors instead of retaining a second complete FC1 payload.
+The original checkpoint is still loaded eagerly; this bounds additional
+interleaving storage, not total checkpoint memory. Direct `preprocess_weights()`
+calls retain eager behavior unless an output `weight_transforms` mapping is
+provided. Pass that mapping to `ModelPackage.apply_weights()` with the returned
+weights to preserve the deferred transformations.
+
 **Export support is not a runtime qualification matrix.** INT3 QMoE nodes
 carry a draft-format/runtime-warning marker; they do not invent an
 approved schema version. Execution of every pair depends on the target ORT

src/mobius/_model_package.py

@@ -43,6 +43,7 @@
 
 from mobius._export_report import ComponentExportReport
 from mobius._optimizations import fold_initializers_after_weights
+from mobius._weight_utils import WeightTransform
 from mobius.adapters import (
     AdapterArtifact,
     AdapterServiceOptions,
@@ -1415,6 +1416,7 @@ def apply_weights(
         prefix_map: dict[str, str] | None = None,
         *,
         fold_constants: bool = True,
+        weight_transforms: dict[str, WeightTransform] | None = None,
     ) -> None:
         """Apply weights from a state dict across component models.
 
@@ -1432,48 +1434,63 @@ def apply_weights(
                 Weights whose name starts with a prefix are applied to the
                 named component (with the prefix stripped). Unmatched weights
                 are applied to all components.
+            weight_transforms: Shape-preserving transforms keyed by names in
+                ``state_dict`` (the preprocessed names, before prefix routing),
+                evaluated without caching during serialization.
         """
         applied: set[str] = set()
 
         if len(self.data) == 1:
             # Single component — apply all weights directly
             model = next(iter(self.data.values()))
-            applied = _apply_weights_to_model(model, state_dict)
+            applied = _apply_weights_to_model(model, state_dict, weight_transforms)
         elif prefix_map is None:
             # No routing info — try every weight against every model
             for model in self.data.values():
-                applied |= _apply_weights_to_model(model, state_dict)
+                applied |= _apply_weights_to_model(model, state_dict, weight_transforms)
         else:
             # Route by prefix
             routed: dict[str, dict[str, torch.Tensor]] = {name: {} for name in self.data}
+            routed_transforms: dict[str, dict[str, WeightTransform]] = {
+                name: {} for name in self.data
+            }
             unmatched: dict[str, torch.Tensor] = {}
-            # Track original HF names for weights that get stripped
-            stripped_to_original: dict[str, str] = {}
+            stripped_to_original: dict[tuple[str, str], str] = {}
 
             for weight_name, tensor in state_dict.items():
                 matched = False
                 for prefix, component in prefix_map.items():
                     if weight_name.startswith(prefix):
                         stripped = weight_name[len(prefix) :].lstrip(".")
                         routed[component][stripped] = tensor
-                        stripped_to_original[stripped] = weight_name
+                        if weight_transforms is not None and weight_name in weight_transforms:
+                            routed_transforms[component][stripped] = weight_transforms[
+                                weight_name
+                            ]
+                        stripped_to_original[(component, stripped)] = weight_name
                         matched = True
                         break
                 if not matched:
                     unmatched[weight_name] = tensor
 
             for component_name, component_weights in routed.items():
                 applied_stripped = _apply_weights_to_model(
-                    self.data[component_name], component_weights
+                    self.data[component_name],
+                    component_weights,
+                    routed_transforms[component_name],
                 )
                 for s in applied_stripped:
-                    applied.add(stripped_to_original.get(s, s))
+                    applied.add(stripped_to_original.get((component_name, s), s))
 
             # Try unmatched weights against all models
             if unmatched:
                 for model in self.data.values():
-                    applied |= _apply_weights_to_model(model, unmatched)
+                    applied |= _apply_weights_to_model(model, unmatched, weight_transforms)
 
+        if weight_transforms is not None:
+            unused = set(weight_transforms) - applied
+            if unused:
+                raise ValueError(f"Weight transforms were not bound: {sorted(unused)!r}.")
         _log_weight_mapping(state_dict, applied)
 
         if not fold_constants:
@@ -1629,7 +1646,11 @@ def _log_weight_mapping(
         )
 
 
-def _apply_weights_to_model(model: ir.Model, state_dict: dict[str, torch.Tensor]) -> set[str]:
+def _apply_weights_to_model(
+    model: ir.Model,
+    state_dict: dict[str, torch.Tensor],
+    weight_transforms: dict[str, WeightTransform] | None = None,
+) -> set[str]:
     """Apply weights to a single model (internal helper).
 
     Uses :func:`_assign_weight` from ``_weight_loading`` for shape
@@ -1643,6 +1664,11 @@ def _apply_weights_to_model(model: ir.Model, state_dict: dict[str, torch.Tensor]
         if name not in model.graph.initializers:
             continue
 
-        _assign_weight(model.graph.initializers[name], tensor, name)
+        _assign_weight(
+            model.graph.initializers[name],
+            tensor,
+            name,
+            transform=weight_transforms.get(name) if weight_transforms is not None else None,
+        )
         applied.add(name)
     return applied

src/mobius/_model_package_test.py

@@ -1440,6 +1440,107 @@ def test_save_preserves_public_policy_symbols(self, tmp_path):
 
 
 class TestModelPackageApplyWeights:
+    @pytest.mark.parametrize("second_has_initializer", [False, True])
+    def test_weight_transform_accounting_separates_component_local_names(
+        self, second_has_initializer
+    ):
+        values = {}
+        models = {}
+        for component in ("text", "vision"):
+            graph = ir.Graph([], [], nodes=[], name=component)
+            if component == "text" or second_has_initializer:
+                value = ir.Value(
+                    name="weight",
+                    type=ir.TensorType(ir.DataType.FLOAT),
+                    shape=ir.Shape([2, 3]),
+                    const_value=ir.tensor(torch.zeros(2, 3)),
+                )
+                graph.register_initializer(value)
+                values[component] = value
+            models[component] = ir.Model(graph, ir_version=10)
+        package = ModelPackage(models)
+        weights = {"text.weight": torch.zeros(2, 3), "vision.weight": torch.ones(2, 3)}
+        transforms = {
+            "text.weight": lambda tensor: tensor + 2,
+            "vision.weight": lambda tensor: tensor + 3,
+        }
+        if not second_has_initializer:
+            with pytest.raises(ValueError, match=r"not bound.*vision\.weight"):
+                package.apply_weights(
+                    weights,
+                    prefix_map={"text.": "text", "vision.": "vision"},
+                    weight_transforms={"vision.weight": transforms["vision.weight"]},
+                    fold_constants=False,
+                )
+        else:
+            package.apply_weights(
+                weights,
+                prefix_map={"text.": "text", "vision.": "vision"},
+                weight_transforms=transforms,
+                fold_constants=False,
+            )
+            for component in ("text", "vision"):
+                actual = torch.from_numpy(values[component].const_value.numpy())
+                expected = transforms[component + ".weight"](weights[component + ".weight"])
+                assert torch.equal(actual, expected)
+
+    @pytest.mark.parametrize("routed", [False, True])
+    def test_weight_transforms_are_lazy_and_follow_prefix_routing(self, routed):
+        value = ir.Value(
+            name="weight",
+            type=ir.TensorType(ir.DataType.FLOAT),
+            shape=ir.Shape([2, 3]),
+            const_value=ir.tensor(torch.zeros(2, 3)),
+        )
+        graph = ir.Graph([], [], nodes=[], name="text")
+        graph.register_initializer(value)
+        package = ModelPackage(
+            {"text": ir.Model(graph, ir_version=10), "other": _make_simple_model()}
+        )
+        calls = []
+        source = torch.arange(6).reshape(2, 3).float()
+
+        def transform(tensor):
+            calls.append(tensor.data_ptr())
+            return tensor + 1
+
+        name = "text.weight" if routed else "weight"
+        package.apply_weights(
+            {name: source},
+            prefix_map={"text.": "text"} if routed else None,
+            weight_transforms={name: transform},
+            fold_constants=False,
+        )
+        assert calls == []
+        assert isinstance(value.const_value, ir.LazyTensor)
+        assert not value.const_value.cache
+        for count in (1, 2):
+            assert torch.equal(torch.from_numpy(value.const_value.numpy()), source + 1)
+            assert len(calls) == count
+
+    def test_unbound_weight_transform_raises(self):
+        package = ModelPackage({"model": _make_simple_model()})
+        with pytest.raises(ValueError, match="Weight transforms were not bound"):
+            package.apply_weights({}, weight_transforms={"missing": lambda tensor: tensor})
+
+    def test_weight_transform_changed_shape_raises(self):
+        value = ir.Value(
+            name="weight",
+            type=ir.TensorType(ir.DataType.FLOAT),
+            shape=ir.Shape([2, 3]),
+            const_value=ir.tensor(torch.zeros(2, 3)),
+        )
+        graph = ir.Graph([], [], nodes=[], name="model")
+        graph.register_initializer(value)
+        package = ModelPackage({"model": ir.Model(graph, ir_version=10)})
+        package.apply_weights(
+            {"weight": torch.zeros(2, 3)},
+            weight_transforms={"weight": lambda tensor: tensor.reshape(3, 2)},
+            fold_constants=False,
+        )
+        with pytest.raises(ValueError, match="changed shape"):
+            value.const_value.numpy()
+
     def test_single_component(self):
         config = make_config()
         module = CausalLMModel(config)