From d9253d450a3b053eda274702f3a79024c4fbfbea Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Wed, 26 Aug 2026 20:13:31 -0700 Subject: [PATCH 1/2] Migrate composite model weight adapters Add an explicit model-weight adapter boundary and migrate Gemma4, Qwen3.5/QMoE, and T5 to preserve architecture-specific semantics while handing component-routed packed weights to the generic loader. Validate every affine quantized op is fully bound. Signed-off-by: Xiaoyu Zhang --- CHANGELOG.md | 3 + src/mobius/__main__.py | 15 +- src/mobius/_component_quantization.py | 108 ++++++++- src/mobius/_component_quantization_test.py | 32 +++ src/mobius/components/_moe.py | 7 +- .../integrations/transformers/_builder.py | 15 +- src/mobius/models/gemma4.py | 185 ++++++++++----- src/mobius/models/gemma4_test.py | 223 +++++++++++++++++- src/mobius/models/qwen35.py | 99 +++++--- src/mobius/models/qwen35_test.py | 29 +++ src/mobius/models/t5.py | 15 ++ src/mobius/weights/__init__.py | 8 + src/mobius/weights/_adapters.py | 68 ++++++ src/mobius/weights/_adapters_test.py | 65 +++++ 14 files changed, 776 insertions(+), 96 deletions(-) create mode 100644 src/mobius/weights/_adapters.py create mode 100644 src/mobius/weights/_adapters_test.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 3a56b6a2e..2d2363bbd 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/src/mobius/__main__.py b/src/mobius/__main__.py index 6d4011acc..72eb475cb 100644 --- a/src/mobius/__main__.py +++ b/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: diff --git a/src/mobius/_component_quantization.py b/src/mobius/_component_quantization.py index 5d7468c47..3483c8d1e 100644 --- a/src/mobius/_component_quantization.py +++ b/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,22 @@ 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_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 +641,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}" + ) diff --git a/src/mobius/_component_quantization_test.py b/src/mobius/_component_quantization_test.py index 09d2b9f90..83e3227fa 100644 --- a/src/mobius/_component_quantization_test.py +++ b/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}, + ), + ) diff --git a/src/mobius/components/_moe.py b/src/mobius/components/_moe.py index 9aa47a249..6714a4dcd 100644 --- a/src/mobius/components/_moe.py +++ b/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 diff --git a/src/mobius/integrations/transformers/_builder.py b/src/mobius/integrations/transformers/_builder.py index fe75adc49..53feafeb8 100644 --- a/src/mobius/integrations/transformers/_builder.py +++ b/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 diff --git a/src/mobius/models/gemma4.py b/src/mobius/models/gemma4.py index df640e3f5..b4e47f0c8 100644 --- a/src/mobius/models/gemma4.py +++ b/src/mobius/models/gemma4.py @@ -33,7 +33,7 @@ from onnxscript import OpBuilder, nn from mobius._build_context import ep_capabilities, is_prefill_prefix_pruning_enabled -from mobius._configs import ArchitectureConfig, Gemma4Config +from mobius._configs import ArchitectureConfig, Gemma4Config, QuantizationConfig from mobius._weight_utils import ( is_packed_quant_key, preprocess_quantized_weights, @@ -51,6 +51,7 @@ ScaleFreeRMSNorm, create_attention_bias, initialize_rope, + make_clippable_quantized_linear_factory, make_quantized_linear_factory, ) from mobius.components._activations import get_activation @@ -139,83 +140,89 @@ def _retain_last_bias_query_row(op: OpBuilder, bias: ir.Value | None) -> ir.Valu return op.Unsqueeze(last, op.Constant(value_ints=[2])) -def _text_quantization_config(config: Gemma4Config): - """Return the active weight-quantization config, or ``None`` when off.""" - quantization_config = getattr(config, "quantization", None) - if quantization_config is None or quantization_config.quant_method == "none": +def _active_quantization( + quantization: QuantizationConfig | None, +) -> QuantizationConfig | None: + if quantization is None or quantization.quant_method == "none": return None - return quantization_config + return quantization -def _quantized_linear_class(config: Gemma4Config) -> type | None: - """Return the checkpoint's QuantizedLinear factory, or ``None`` when off.""" - quantization_config = _text_quantization_config(config) - if quantization_config is None: +def _component_quantization_config( + config: Gemma4Config, + component: str, + *, + source_module_names: tuple[str, ...] = (), +) -> QuantizationConfig | None: + if config.component_quantization is not None: + quantization = config.quantization_for(component) + else: + quantization = config.quantization + if component == "vision_encoder" and ( + quantization is None or not quantization.quantize_vision + ): + return None + if component == "audio_encoder": + return None + quantization = _active_quantization(quantization) + if quantization is not None and source_module_names: + quantization = quantization.for_module(source_module_names) + return quantization + + +def _quantized_linear_class( + config: Gemma4Config, + quantization: QuantizationConfig | None, +) -> type | None: + if quantization is None: return None zero_point_dtype = ( - config.dtype - if getattr(quantization_config, "float_zero_point", False) - else ir.DataType.UINT8 + config.dtype if getattr(quantization, "float_zero_point", False) else ir.DataType.UINT8 ) return make_quantized_linear_factory( - bits=quantization_config.bits, - block_size=quantization_config.group_size, - has_zero_point=not quantization_config.sym, + bits=quantization.bits, + block_size=quantization.group_size, + has_zero_point=not quantization.sym, zero_point_dtype=zero_point_dtype, ) def _text_linear_class(config: Gemma4Config) -> type | None: """Return a QuantizedLinear factory for text projections, or ``None``.""" - return _quantized_linear_class(config) + return _quantized_linear_class( + config, + _component_quantization_config(config, "decoder"), + ) def _vision_linear_classes(config: Gemma4Config) -> tuple[type, type]: """Return plain and activation-clipped Linear classes for the vision graph.""" - quantization_config = _text_quantization_config(config) - quantized_linear = _quantized_linear_class(config) - if ( - quantization_config is None - or not quantization_config.quantize_vision - or quantized_linear is None - ): + quantization = _component_quantization_config(config, "vision_encoder") + quantized_linear = _quantized_linear_class(config, quantization) + if quantization is None or quantized_linear is None: return Linear, ClippableLinear - - class QuantizedClippableLinear(quantized_linear): - """MatMulNBits projection with Gemma4's learned activation clipping.""" - - def __init__(self, in_features: int, out_features: int, bias: bool = False): - super().__init__(in_features, out_features, bias=bias) - self.input_min = nn.Parameter([]) - self.input_max = nn.Parameter([]) - self.output_min = nn.Parameter([]) - self.output_max = nn.Parameter([]) - - def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value: - x = op.Clip(x, self.input_min, self.input_max) - return op.Clip( - super().forward(op, x), - self.output_min, - self.output_max, - ) - - return quantized_linear, QuantizedClippableLinear - - -def _text_embeddings_quantized(config: Gemma4Config) -> bool: - """Whether the text token-embedding tables use GatherBlockQuantized.""" - quantization_config = _text_quantization_config(config) - return quantization_config is not None and bool( - getattr(quantization_config, "quantize_embeddings", False) + zero_point_dtype = ( + config.dtype if getattr(quantization, "float_zero_point", False) else ir.DataType.UINT8 + ) + return ( + quantized_linear, + make_clippable_quantized_linear_factory( + bits=quantization.bits, + block_size=quantization.group_size, + has_zero_point=not quantization.sym, + zero_point_dtype=zero_point_dtype, + ), ) def _text_lm_head_quantized(config: Gemma4Config) -> bool: """Whether the text LM head projection uses MatMulNBits.""" - quantization_config = _text_quantization_config(config) - return quantization_config is not None and bool( - getattr(quantization_config, "quantize_lm_head", False) + quantization = _component_quantization_config( + config, + "decoder", + source_module_names=("lm_head", "model.language_model.lm_head"), ) + return quantization is not None and bool(getattr(quantization, "quantize_lm_head", False)) def _make_scaled_word_embedding( @@ -223,6 +230,9 @@ def _make_scaled_word_embedding( num_embeddings: int, embedding_dim: int, embed_scale: float, + *, + component: str = "decoder", + source_module_names: tuple[str, ...] = ("model.language_model.embed_tokens",), ): """Build a scaled token embedding, quantized when the config requests it. @@ -230,9 +240,14 @@ def _make_scaled_word_embedding( lookup) when embedding quantization is enabled and the embedding dimension is block-aligned, otherwise a float :class:`Gemma3TextScaledWordEmbedding`. """ - quantization_config = _text_quantization_config(config) + quantization_config = _component_quantization_config( + config, + component, + source_module_names=source_module_names, + ) if ( - _text_embeddings_quantized(config) + quantization_config is not None + and quantization_config.quantize_embeddings and embedding_dim % quantization_config.group_size == 0 ): return Gemma4ScaledQuantizedWordEmbedding( @@ -1952,6 +1967,7 @@ def __init__(self, config: Gemma4Config): vocab_per_layer, self._num_layers * self._per_layer_dim, float(self._per_layer_dim**0.5), + source_module_names=("model.language_model.embed_tokens_per_layer",), ) # Split [V, D] tables — used when split_per_layer_embedding is True # (i.e. the fused table exceeds the EP's max_buffer_size, e.g. WebGPU's @@ -2767,6 +2783,8 @@ def __init__(self, config: Gemma4Config): config.vocab_size, config.hidden_size, embed_scale, + component="embedding", + source_module_names=("model.language_model.embed_tokens",), ) self.image_token_id = config.image_token_id or 0 # Audio token ID is only set when the model has an audio encoder. @@ -2786,6 +2804,8 @@ def __init__(self, config: Gemma4Config): vocab_per_layer, self._num_layers * self._per_layer_dim, float(self._per_layer_dim**0.5), + component="embedding", + source_module_names=("model.language_model.embed_tokens_per_layer",), ) self.per_layer_model_projection = Linear( config.hidden_size, @@ -3305,6 +3325,25 @@ class Gemma4Model(nn.Module): "model.language_model.per_layer_projection_norm", ), } + HF_COMPONENT_MODULE_ALIASES: ClassVar[dict[str, dict[str, str]]] = { + "decoder": { + "model": "model.language_model", + "lm_head": "lm_head", + }, + "vision_encoder": { + "encoder": "model.vision_tower.encoder", + "projector": "model.embed_vision.embedding_projection", + }, + "audio_encoder": { + "encoder": "model.audio_tower", + "projector": "model.embed_audio.embedding_projection", + }, + "embedding": { + "embed_tokens": "model.language_model.embed_tokens", + "embed_tokens_per_layer": ("model.language_model.embed_tokens_per_layer"), + "per_layer_model_projection": ("model.language_model.per_layer_model_projection"), + }, + } def __init__(self, config: Gemma4Config): super().__init__() @@ -3384,6 +3423,14 @@ def preprocess_weights( elif any(suffix.startswith(p) for p in per_layer_prefixes): # Per-layer embedding weights → embedding sub-model renamed["embedding." + suffix] = value + elif ( + suffix.startswith("embed_tokens.") + and self.config.component_quantization is not None + ): + # The split decoder consumes inputs_embeds and has no token + # table initializer. Route authoritative component-plan + # embedding sidecars only to the embedding graph. + renamed["embedding." + suffix] = value else: # All other text weights nest under decoder.model.* onnx_key = "decoder.model." + suffix @@ -3461,11 +3508,11 @@ def preprocess_weights( _remap_moe_expert_weights(renamed, self.config) quantization = self.config.quantization - if quantization is not None and quantization.quant_method in { - "olive", - "gptq", - "awq", - }: + if ( + self.config.component_quantization is None + and quantization is not None + and quantization.quant_method in {"olive", "gptq", "awq"} + ): tie = self.config.tie_word_embeddings apply_tie = tie and any( key in renamed @@ -3559,6 +3606,22 @@ class Gemma4UnifiedModel(nn.Module): "audio_encoder": ("model.embed_audio",), "embedding": ("model.language_model.embed_tokens",), } + HF_COMPONENT_MODULE_ALIASES: ClassVar[dict[str, dict[str, str]]] = { + "decoder": { + "model": "model.language_model", + "lm_head": "lm_head", + }, + "vision_encoder": { + "patch_dense": "model.vision_embedder.patch_dense", + "projector": "model.embed_vision.embedding_projection", + }, + "audio_encoder": { + "projector": "model.embed_audio.embedding_projection", + }, + "embedding": { + "embed_tokens": "model.language_model.embed_tokens", + }, + } def __init__(self, config: Gemma4Config): super().__init__() diff --git a/src/mobius/models/gemma4_test.py b/src/mobius/models/gemma4_test.py index 8e8c90bd6..8eb696003 100644 --- a/src/mobius/models/gemma4_test.py +++ b/src/mobius/models/gemma4_test.py @@ -11,7 +11,12 @@ import pytest import torch -from mobius._configs import AudioConfig, Gemma4Config, QuantizationConfig +from mobius._configs import ( + AudioConfig, + Gemma4AudioConfig, + Gemma4Config, + QuantizationConfig, +) from mobius.models.gemma4 import Gemma4CausalLMModel, Gemma4EmbeddingModel, Gemma4Model @@ -382,6 +387,222 @@ def test_component_regex_keeps_per_layer_decoder_projections_float(): assert type(layer.per_layer_projection) is Linear +class TestGemma4ComponentWeightAdapters: + @staticmethod + def _config() -> Gemma4Config: + decoder = QuantizationConfig( + bits=4, + group_size=16, + quant_method="olive", + sym=True, + modules_to_not_convert=( + r"re:.*\.per_layer_input_gate", + r"re:.*\.per_layer_projection", + ), + ) + return _tiny_gemma4_config( + enable_moe_block=False, + hidden_size_per_layer_input=16, + vocab_size_per_layer_input=256, + audio=Gemma4AudioConfig( + input_size=32, + num_layers=1, + hidden_size=32, + output_proj_dims=64, + subsampling_conv_channels=[16, 8], + audio_token_id=254, + ), + quantization=decoder, + component_quantization={ + "decoder": decoder, + "vision_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, + ), + "embedding": QuantizationConfig( + bits=8, + group_size=16, + quant_method="olive", + sym=True, + ), + }, + ) + + @staticmethod + def _layouts(graph) -> set[tuple[int, int]]: + return { + ( + node.attributes["bits"].as_int(), + node.attributes["block_size"].as_int(), + ) + for node in graph + if node.op_type == "MatMulNBits" + } + + def test_builds_each_component_with_its_declared_layout(self): + from mobius._builder import build_from_module + + config = self._config() + package = build_from_module( + Gemma4Model(config), + config, + task="gemma4", + ) + + assert self._layouts(package["decoder"].graph) == {(4, 16)} + assert self._layouts(package["vision_encoder"].graph) == {(8, 32)} + assert self._layouts(package["audio_encoder"].graph) == {(2, 16)} + assert self._layouts(package["embedding"].graph) == {(8, 16)} + + def test_vision_full_hf_linear_exclusion_matches_local_module(self): + from mobius._component_quantization import configure_component_quantization + from mobius.components import ClippableLinear, ClippableQuantizedLinear + from mobius.tasks._gemma4 import Gemma4Task + + config = self._config() + vision_quantization = dataclasses.replace( + config.component_quantization["vision_encoder"], + modules_to_not_convert=( + "model.vision_tower.encoder.layers.0.self_attn.q_proj.linear", + ), + ) + config = dataclasses.replace( + config, + vision=dataclasses.replace( + config.vision, + use_clipped_linears=True, + ), + component_quantization={ + **config.component_quantization, + "vision_encoder": vision_quantization, + }, + ) + module = Gemma4Model(config) + + configure_component_quantization(module, config, Gemma4Task()) + + attention = module.vision_encoder.encoder.layers[0].self_attn + assert type(attention.q_proj) is ClippableLinear + assert isinstance(attention.k_proj, ClippableQuantizedLinear) + + def test_routes_raw_sidecars_to_each_component_codec(self): + from mobius._component_quantization import ( + configure_component_quantization, + normalize_component_quantized_weights, + ) + from mobius.tasks._gemma4 import Gemma4Task + + config = self._config() + module = Gemma4Model(config) + task = Gemma4Task() + manifest = configure_component_quantization(module, config, task) + state_dict = { + "model.language_model.layers.0.self_attn.q_proj.weight_qweight": torch.zeros( + 64, 32, dtype=torch.uint8 + ), + "model.language_model.layers.0.self_attn.q_proj.weight_scales": torch.ones(64, 4), + "model.vision_tower.encoder.layers.0.self_attn.q_proj.linear.weight_qweight": torch.zeros( + 32, 32, dtype=torch.uint8 + ), + "model.vision_tower.encoder.layers.0.self_attn.q_proj.linear.weight_scales": torch.ones( + 32, 1 + ), + "model.audio_tower.layers.0.self_attn.q_proj.linear.weight_qweight": torch.zeros( + 32, 8, dtype=torch.uint8 + ), + "model.audio_tower.layers.0.self_attn.q_proj.linear.weight_scales": torch.ones( + 32, 2 + ), + "model.language_model.per_layer_model_projection.weight_qweight": torch.zeros( + 32, 64, dtype=torch.uint8 + ), + "model.language_model.per_layer_model_projection.weight_scales": torch.ones(32, 4), + } + + renamed = module.preprocess_weights(state_dict) + result = normalize_component_quantized_weights( + renamed, + module, + config, + ("decoder", "vision_encoder", "audio_encoder", "embedding"), + manifest=manifest, + task=task, + ) + + assert result["decoder.model.layers.0.self_attn.q_proj.weight"].shape == ( + 64, + 4, + 8, + ) + assert result["vision_encoder.encoder.layers.0.self_attn.q_proj.weight"].shape == ( + 32, + 1, + 32, + ) + assert result["audio_encoder.encoder.layers.0.self_attn.q_proj.weight"].shape == ( + 32, + 2, + 4, + ) + assert result["embedding.per_layer_model_projection.weight"].shape == ( + 32, + 4, + 16, + ) + + def test_quantized_embedding_sidecars_route_only_to_embedding_component(self): + from mobius._component_quantization import ( + configure_component_quantization, + normalize_component_quantized_weights, + ) + from mobius.tasks._gemma4 import Gemma4Task + + config = self._config() + embedding_quantization = dataclasses.replace( + config.component_quantization["embedding"], + quantize_embeddings=True, + ) + config = dataclasses.replace( + config, + component_quantization={ + **config.component_quantization, + "embedding": embedding_quantization, + }, + ) + module = Gemma4Model(config) + task = Gemma4Task() + manifest = configure_component_quantization(module, config, task) + renamed = module.preprocess_weights( + { + "model.language_model.embed_tokens.weight_qweight": torch.zeros( + 256, 64, dtype=torch.uint8 + ), + "model.language_model.embed_tokens.weight_scales": torch.ones(256, 4), + } + ) + + assert not any(key.startswith("decoder.") for key in renamed) + result = normalize_component_quantized_weights( + renamed, + module, + config, + ("decoder", "vision_encoder", "audio_encoder", "embedding"), + manifest=manifest, + task=task, + ) + + assert result["embedding.embed_tokens.qweight"].shape == (256, 64) + assert result["embedding.embed_tokens.scales"].shape == (256, 4) + + class TestScaleFreeRMSNormOverflow: """V norm should handle FP16 overflow from squaring large values.""" diff --git a/src/mobius/models/qwen35.py b/src/mobius/models/qwen35.py index 96b12961a..a62fe2ebd 100644 --- a/src/mobius/models/qwen35.py +++ b/src/mobius/models/qwen35.py @@ -9,7 +9,7 @@ import torch from onnxscript import OpBuilder, nn -from mobius._configs import ArchitectureConfig +from mobius._configs import ArchitectureConfig, QuantizationConfig from mobius._weight_utils import ( preprocess_quantized_weights, supported_qmoe_quantization, @@ -42,6 +42,14 @@ # --------------------------------------------------------------------------- +def _decoder_quantization( + config: ArchitectureConfig, +) -> QuantizationConfig | None: + if config.component_quantization is not None: + return config.quantization_for("decoder") + return config.quantization + + def _linear_factory(config: ArchitectureConfig) -> type | None: """Build a quantized-linear factory from ``config.quantization``, or None. @@ -53,7 +61,7 @@ def _linear_factory(config: ArchitectureConfig) -> type | None: A few modules opt out of this factory for specific quantizers — see :data:`_FLOAT_MODULE_QUANT_METHODS`. """ - quantization = config.quantization + quantization = _decoder_quantization(config) if quantization is None or quantization.quant_method == "none": return None zero_point_dtype = config.dtype if quantization.float_zero_point else ir.DataType.UINT8 @@ -85,7 +93,7 @@ def _linear_factory(config: ArchitectureConfig) -> type | None: def _keeps_modules_float(config: ArchitectureConfig) -> bool: """True when the checkpoint's quantizer leaves the opt-out modules float.""" - quantization = config.quantization + quantization = _decoder_quantization(config) return ( quantization is not None and quantization.quant_method in _FLOAT_MODULE_QUANT_METHODS ) @@ -149,7 +157,7 @@ def _linear_attn_class( ``MatMulNBits`` initializers the checkpoint never contains); otherwise the same factory used for the rest of the layer. """ - quantization = config.quantization + quantization = _decoder_quantization(config) method = quantization.quant_method if quantization is not None else None if method in cls.float_linear_attn_quant_methods: return None @@ -504,7 +512,7 @@ def preprocess_weights( # fused expert-major tensors and route them through the QMoE repacker # instead of un-fusing into per-expert MLPs. Uses the same predicate # as MoELayer so the weights and the emitted graph never disagree. - use_qmoe = supported_qmoe_quantization(self.config.quantization) is not None + use_qmoe = supported_qmoe_quantization(_decoder_quantization(self.config)) is not None cleaned: dict[str, torch.Tensor] = {} for key, value in state_dict.items(): if key.startswith(("mtp_", "mtp.")): @@ -554,7 +562,7 @@ def preprocess_weights( return preprocess_quantized_weights( cleaned, - self.config.quantization, + _decoder_quantization(self.config), tie_embeddings=effective_tie_word_embeddings(self.config), qmoe_target_path=".mlp", qmoe_quant_methods=("gptq", "awq", "olive"), @@ -652,7 +660,7 @@ def preprocess_weights( elif stripped.startswith("language_model."): suffix = stripped[len("language_model.") :] renamed[f"decoder.model.{suffix}"] = value - quantization = self.config.quantization + quantization = _decoder_quantization(self.config) tie = effective_tie_word_embeddings(self.config) # Preserve the old VL partial-state-dict behavior: tying is a no-op # when neither decoder table is present. Production builds pass the @@ -661,15 +669,33 @@ def preprocess_weights( key in renamed for key in ("decoder.model.embed_tokens.weight", "decoder.lm_head.weight") ) - result = preprocess_quantized_weights( - renamed, - quantization, - tie_embeddings=apply_tie, - embed_key="decoder.model.embed_tokens.weight", - head_key="decoder.lm_head.weight", - qmoe_target_path=None, - reject_quantized_embeddings_lm_head=True, - ) + if self.config.component_quantization is not None: + decoder_weights = { + key: value for key, value in renamed.items() if key.startswith("decoder.") + } + other_weights = { + key: value for key, value in renamed.items() if not key.startswith("decoder.") + } + result = preprocess_quantized_weights( + decoder_weights, + self.config.quantization_for("decoder"), + tie_embeddings=apply_tie, + embed_key="decoder.model.embed_tokens.weight", + head_key="decoder.lm_head.weight", + qmoe_target_path=None, + reject_quantized_embeddings_lm_head=True, + ) + result.update(other_weights) + else: + result = preprocess_quantized_weights( + renamed, + quantization, + tie_embeddings=apply_tie, + embed_key="decoder.model.embed_tokens.weight", + head_key="decoder.lm_head.weight", + qmoe_target_path=None, + reject_quantized_embeddings_lm_head=True, + ) if tie: if ( "decoder.model.embed_tokens.weight" not in result @@ -854,7 +880,7 @@ def preprocess_weights( native QMoE ABI, they instead remain expert-major, are renamed from Olive's suffix convention, and are packed into QMoE parameters. """ - quantization = self.config.quantization + quantization = _decoder_quantization(self.config) use_qmoe = supported_qmoe_quantization(quantization) is not None tie = effective_tie_word_embeddings(self.config) @@ -906,16 +932,35 @@ def preprocess_weights( key in renamed for key in ("decoder.model.embed_tokens.weight", "decoder.lm_head.weight") ) - result = preprocess_quantized_weights( - renamed, - quantization, - tie_embeddings=apply_tie, - embed_key="decoder.model.embed_tokens.weight", - head_key="decoder.lm_head.weight", - qmoe_target_path=".mlp", - qmoe_quant_methods=("olive",), - reject_quantized_embeddings_lm_head=True, - ) + if self.config.component_quantization is not None: + decoder_weights = { + key: value for key, value in renamed.items() if key.startswith("decoder.") + } + other_weights = { + key: value for key, value in renamed.items() if not key.startswith("decoder.") + } + result = preprocess_quantized_weights( + decoder_weights, + self.config.quantization_for("decoder"), + tie_embeddings=apply_tie, + embed_key="decoder.model.embed_tokens.weight", + head_key="decoder.lm_head.weight", + qmoe_target_path=".mlp", + qmoe_quant_methods=("olive",), + reject_quantized_embeddings_lm_head=True, + ) + result.update(other_weights) + else: + result = preprocess_quantized_weights( + renamed, + quantization, + tie_embeddings=apply_tie, + embed_key="decoder.model.embed_tokens.weight", + head_key="decoder.lm_head.weight", + qmoe_target_path=".mlp", + qmoe_quant_methods=("olive",), + reject_quantized_embeddings_lm_head=True, + ) if tie: if ( "decoder.model.embed_tokens.weight" not in result diff --git a/src/mobius/models/qwen35_test.py b/src/mobius/models/qwen35_test.py index efcf60026..1756f0688 100644 --- a/src/mobius/models/qwen35_test.py +++ b/src/mobius/models/qwen35_test.py @@ -480,6 +480,35 @@ def test_unsupported_qmoe_quantization_with_packed_experts_raises(self): class TestQwen35MoEVL3ModelQMoEExport: + def test_plan_only_quantization_keeps_graph_and_qmoe_weights_aligned(self): + quantization = QuantizationConfig( + bits=4, + group_size=_BLK, + quant_method="olive", + sym=False, + ) + config = dataclasses.replace( + _moe_vl_config(None), + component_quantization={"decoder": quantization}, + ) + model = Qwen35MoEVL3ModelCausalLMModel(config) + + block = model.decoder.model.layers[0].mlp + assert block.experts is None + assert hasattr(block, "fc1_experts_weights") + + result = model.preprocess_weights(_olive_expert_state_dict()) + prefix = "decoder.model.layers.0.mlp." + parameter_names = {name for name, _ in model.named_parameters()} + for suffix in ( + "fc1_experts_weights", + "fc1_scales", + "fc2_experts_weights", + "fc2_scales", + ): + assert prefix + suffix in result + assert prefix + suffix in parameter_names + def test_olive_preprocess_packs_qmoe_and_binds(self): config = _moe_vl_config( QuantizationConfig(bits=4, group_size=_BLK, quant_method="olive", sym=False) diff --git a/src/mobius/models/t5.py b/src/mobius/models/t5.py index 5bc191537..6f8e5c6b6 100644 --- a/src/mobius/models/t5.py +++ b/src/mobius/models/t5.py @@ -14,6 +14,7 @@ from mobius._configs import ArchitectureConfig from mobius._weight_utils import preprocess_quantized_weights +from mobius._weight_utils import tie_word_embeddings as tie_weight_tensors from mobius.components._activations import ACT2FN from mobius.components._common import Embedding, Linear from mobius.components._encoder_decoder_attention import ( @@ -554,6 +555,20 @@ def preprocess_weights( embed = new_state_dict.get("encoder.embed_tokens.weight") if embed is not None: new_state_dict["decoder.lm_head.weight"] = embed + if self.config.component_quantization is not None: + tie_weight_tensors( + new_state_dict, + embed_key="encoder.embed_tokens.weight", + head_key="decoder.lm_head.weight", + ) + if ( + "encoder.embed_tokens.weight" in new_state_dict + and "decoder.embed_tokens.weight" not in new_state_dict + ): + new_state_dict["decoder.embed_tokens.weight"] = new_state_dict[ + "encoder.embed_tokens.weight" + ] + return new_state_dict return preprocess_quantized_weights( new_state_dict, self.config.quantization, diff --git a/src/mobius/weights/__init__.py b/src/mobius/weights/__init__.py index a3dd7744d..7f6d9594f 100644 --- a/src/mobius/weights/__init__.py +++ b/src/mobius/weights/__init__.py @@ -5,6 +5,11 @@ from __future__ import annotations +from mobius.weights._adapters import ( + ModelWeightAdapter, + WeightAdapterContext, + adapt_model_weights, +) from mobius.weights._codecs import ( QuantizationCodec, QuantizationCodecRegistry, @@ -19,10 +24,13 @@ __all__ = [ "FloatWeight", + "ModelWeightAdapter", "PackedWeight", "QuantizationCodec", "QuantizationCodecRegistry", "WeightBundle", + "WeightAdapterContext", "WeightRecord", "codec_registry", + "adapt_model_weights", ] diff --git a/src/mobius/weights/_adapters.py b/src/mobius/weights/_adapters.py new file mode 100644 index 000000000..fcb6f396e --- /dev/null +++ b/src/mobius/weights/_adapters.py @@ -0,0 +1,68 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Narrow model-specific boundary in the checkpoint loading pipeline.""" + +from __future__ import annotations + +__all__ = [ + "ModelWeightAdapter", + "WeightAdapterContext", + "adapt_model_weights", +] + +import dataclasses +from collections.abc import Mapping +from typing import Any, Protocol + +import torch +from onnxscript import nn + +from mobius._component_manifest import ComponentManifest +from mobius._configs import BaseModelConfig + + +@dataclasses.dataclass(frozen=True) +class WeightAdapterContext: + """Generic metadata available to a model-specific semantic adapter.""" + + config: BaseModelConfig + manifest: ComponentManifest + + +class ModelWeightAdapter(Protocol): + """Architecture-specific rename/split/fuse operations only.""" + + def adapt( + self, + module: nn.Module, + state_dict: Mapping[str, torch.Tensor], + context: WeightAdapterContext, + ) -> dict[str, torch.Tensor]: + """Return semantically aligned weights without format normalization.""" + ... + + +def adapt_model_weights( + module: nn.Module, + state_dict: Mapping[str, torch.Tensor], + *, + config: BaseModelConfig, + manifest: ComponentManifest, +) -> dict[str, torch.Tensor]: + """Run an explicit adapter or the legacy ``preprocess_weights`` hook.""" + context = WeightAdapterContext(config=config, manifest=manifest) + adapter: ModelWeightAdapter | None = getattr(module, "weight_adapter", None) + if adapter is not None: + return adapter.adapt(module, state_dict, context) + + preprocess = getattr(module, "preprocess_weights", None) + if preprocess is None: + return dict(state_dict) + result: Any = preprocess(dict(state_dict)) + if not isinstance(result, dict): + raise TypeError( + f"{type(module).__name__}.preprocess_weights must return a dict, " + f"got {type(result).__name__}" + ) + return result diff --git a/src/mobius/weights/_adapters_test.py b/src/mobius/weights/_adapters_test.py new file mode 100644 index 000000000..d26499795 --- /dev/null +++ b/src/mobius/weights/_adapters_test.py @@ -0,0 +1,65 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Tests for the model weight adapter boundary.""" + +from __future__ import annotations + +import torch +from onnxscript import nn + +from mobius._component_manifest import ComponentDescriptor, ComponentManifest +from mobius._configs import ArchitectureConfig +from mobius.weights import WeightAdapterContext, adapt_model_weights + + +def _manifest() -> ComponentManifest: + return ComponentManifest( + ( + ComponentDescriptor( + name="model", + module_path="", + role="decoder", + ), + ) + ) + + +class _Module(nn.Module): + def preprocess_weights(self, state_dict): + return {f"legacy.{name}": value for name, value in state_dict.items()} + + +def test_legacy_preprocess_hook_remains_supported(): + tensor = torch.ones(2) + + result = adapt_model_weights( + _Module(), + {"weight": tensor}, + config=ArchitectureConfig(), + manifest=_manifest(), + ) + + assert result["legacy.weight"] is tensor + + +def test_explicit_adapter_takes_precedence(): + class _Adapter: + def adapt(self, module, state_dict, context: WeightAdapterContext): + assert isinstance(module, _Module) + assert context.manifest.names == ("model",) + return {f"adapter.{name}": value for name, value in state_dict.items()} + + module = _Module() + module.weight_adapter = _Adapter() + tensor = torch.ones(2) + + result = adapt_model_weights( + module, + {"weight": tensor}, + config=ArchitectureConfig(), + manifest=_manifest(), + ) + + assert result["adapter.weight"] is tensor + assert "legacy.weight" not in result From f998a0bb5f9b8808fbd7db4e6d9e9cf36a9caaf7 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Thu, 27 Aug 2026 14:10:09 -0700 Subject: [PATCH 2/2] Use clarified manifest paths in adapters Update binding validation and adapter fixtures for ComponentDescriptor.module_attribute_path. Signed-off-by: Xiaoyu Zhang --- src/mobius/_component_quantization.py | 5 ++++- src/mobius/weights/_adapters_test.py | 2 +- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/src/mobius/_component_quantization.py b/src/mobius/_component_quantization.py index 3483c8d1e..bc159e0df 100644 --- a/src/mobius/_component_quantization.py +++ b/src/mobius/_component_quantization.py @@ -599,7 +599,10 @@ 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_path) + component_module = _resolve_module( + module, + descriptor.module_attribute_path, + ) local_module = ( _resolve_module(component_module, local_path) if component_module is not None diff --git a/src/mobius/weights/_adapters_test.py b/src/mobius/weights/_adapters_test.py index d26499795..36a277269 100644 --- a/src/mobius/weights/_adapters_test.py +++ b/src/mobius/weights/_adapters_test.py @@ -18,7 +18,7 @@ def _manifest() -> ComponentManifest: ( ComponentDescriptor( name="model", - module_path="", + module_attribute_path="", role="decoder", ), )