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 01cce5004..0e64fa4a7 100644 --- a/src/mobius/__main__.py +++ b/src/mobius/__main__.py @@ -549,6 +549,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, + ) + _reject_unsupported_affine_qwen4(model_type, config) if compressed_tensors_config is not None: # Packed FP4 weights cannot pass through ordinary apply_weights. @@ -573,8 +577,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, ) @@ -588,6 +598,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_manifest.py b/src/mobius/_component_manifest.py index 61babcfb7..070c8661a 100644 --- a/src/mobius/_component_manifest.py +++ b/src/mobius/_component_manifest.py @@ -158,10 +158,10 @@ def resolve_component_manifest( component_sources: dict[str, tuple[str, ...]] = {} component_aliases: dict[str, tuple[tuple[str, str], ...]] = {} - if module_class is not None and model_type is not None and hf_config is not None: + if module_class is not None and hf_config is not None: component_sources = get_hf_component_sources( module_class, - model_type, + model_type or "", hf_config, ) raw_aliases = getattr(module_class, "HF_COMPONENT_MODULE_ALIASES", {}) diff --git a/src/mobius/_component_quantization.py b/src/mobius/_component_quantization.py index 40ab5c471..2bedd2edf 100644 --- a/src/mobius/_component_quantization.py +++ b/src/mobius/_component_quantization.py @@ -9,6 +9,7 @@ "attach_hf_component_sources", "configure_component_quantization", "normalize_component_quantized_weights", + "validate_quantized_component_bindings", "preprocess_component_quantized_state_dict", ] @@ -129,17 +130,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, @@ -161,6 +208,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: @@ -192,6 +244,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 " @@ -211,6 +268,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__}); " @@ -539,10 +601,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( @@ -571,6 +648,38 @@ def normalize_component_quantized_weights( 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}" + ) + + def attach_hf_component_sources( module: nn.Module, *, 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 005547f1d..fc5580104 100644 --- a/src/mobius/integrations/transformers/_builder.py +++ b/src/mobius/integrations/transformers/_builder.py @@ -16,6 +16,7 @@ from mobius._component_quantization import ( attach_hf_component_sources, normalize_component_quantized_weights, + validate_quantized_component_bindings, ) from mobius._model_package import ModelPackage from mobius._registry import registry @@ -29,6 +30,7 @@ stream_compressed_tensors_to_package, ) from mobius.tasks import ModelTask +from mobius.weights import adapt_model_weights logger = logging.getLogger(__name__) @@ -594,8 +596,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, @@ -608,6 +614,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 19e147ea6..68b68eb6b 100644 --- a/src/mobius/models/gemma4.py +++ b/src/mobius/models/gemma4.py @@ -165,12 +165,14 @@ def _active_quantization( def _component_quantization_config( config: Gemma4Config, component: str, + *, + source_module_names: tuple[str, ...] = (), ) -> QuantizationConfig | None: """Return the effective packed-linear layout for one Gemma4 component.""" if config.component_quantization is not None: quantization = config.quantization_for_source_paths( component, - _GEMMA4_COMPONENT_SOURCES.get(component, ()), + source_module_names or _GEMMA4_COMPONENT_SOURCES.get(component, ()), ) return _active_quantization(quantization) @@ -191,12 +193,18 @@ def _component_quantization_config( def _table_quantization_config( config: Gemma4Config, component: str, + *, + source_module_names: tuple[str, ...] = (), ) -> QuantizationConfig | None: """Return the config controlling embedding tables in *component*.""" if config.component_quantization is None: return _active_quantization(config.quantization) - return _active_quantization(config.quantization_for(component)) - + return _active_quantization( + config.quantization_for_source_paths( + component, + source_module_names or _GEMMA4_COMPONENT_SOURCES.get(component, ()), + ) + ) def _quantized_linear_class( config: Gemma4Config, @@ -285,6 +293,7 @@ def _make_scaled_word_embedding( 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. @@ -292,7 +301,11 @@ 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 = _table_quantization_config(config, component) + quantization_config = _table_quantization_config( + config, + component, + source_module_names=source_module_names, + ) if ( quantization_config is not None and quantization_config.quantize_embeddings @@ -437,7 +450,7 @@ def _preprocess_component_quantized_weights( embed_key="embedding.embed_tokens.weight", head_key="decoder.lm_head.weight", qmoe_target_path=None, - reject_quantized_embeddings_lm_head=True, + reject_quantized_embeddings_lm_head=component != "embedding", ) ) return result @@ -2128,6 +2141,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 @@ -3471,6 +3485,25 @@ class Gemma4Model(nn.Module): # Runtime HF ``named_modules()`` sub-trees per ONNX component. HF_COMPONENT_SOURCES: ClassVar[dict[str, tuple[str, ...]]] = _GEMMA4_COMPONENT_SOURCES + 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__() @@ -3554,6 +3587,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 @@ -3705,6 +3746,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 5338d8a62..eddf9c3fd 100644 --- a/src/mobius/models/gemma4_test.py +++ b/src/mobius/models/gemma4_test.py @@ -619,6 +619,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 62647cbb8..6cd522dca 100644 --- a/src/mobius/models/qwen35.py +++ b/src/mobius/models/qwen35.py @@ -10,7 +10,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, @@ -44,6 +44,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. @@ -55,7 +63,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 @@ -93,12 +101,12 @@ def _decoder_component_config( if config.component_quantization is None: return config quantization = config.quantization_for_source_paths("decoder", source_paths) - return dataclasses.replace(config, quantization=quantization) + return dataclasses.replace(config, quantization=quantization, component_quantization=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 ) @@ -165,7 +173,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 @@ -523,7 +531,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.")): @@ -573,7 +581,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"), @@ -614,6 +622,7 @@ class Qwen35VL3ModelCausalLMModel(nn.Module): # Runtime HF ``named_modules()`` sub-trees per ONNX component. HF_COMPONENT_SOURCES: ClassVar[dict[str, tuple[str, ...]]] = { "decoder": ( + "model.language_model", "model.language_model.layers", "model.language_model.norm", "model.language_model.rotary_emb", @@ -681,7 +690,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 @@ -868,6 +877,7 @@ class Qwen35MoEVL3ModelCausalLMModel(nn.Module): # Runtime HF ``named_modules()`` sub-trees per ONNX component. HF_COMPONENT_SOURCES: ClassVar[dict[str, tuple[str, ...]]] = { "decoder": ( + "model.language_model", "model.language_model.layers", "model.language_model.norm", "model.language_model.rotary_emb", diff --git a/src/mobius/models/qwen35_test.py b/src/mobius/models/qwen35_test.py index c8ab5c05b..8366ab026 100644 --- a/src/mobius/models/qwen35_test.py +++ b/src/mobius/models/qwen35_test.py @@ -481,6 +481,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 2baa0ee6b..cc7608ff6 100644 --- a/src/mobius/models/t5.py +++ b/src/mobius/models/t5.py @@ -5,6 +5,7 @@ from __future__ import annotations +import dataclasses import math from typing import ClassVar @@ -540,8 +541,27 @@ class T5ForConditionalGeneration(nn.Module): def __init__(self, config: ArchitectureConfig): super().__init__() self.config = config - self.encoder = T5Encoder(config) - self.decoder = T5Decoder(config) + encoder_config = config + decoder_config = config + if config.component_quantization is not None: + encoder_config = dataclasses.replace( + config, + quantization=config.quantization_for_source_paths( + "encoder", + self.HF_COMPONENT_SOURCES["encoder"], + ), + component_quantization=None, + ) + decoder_config = dataclasses.replace( + config, + quantization=config.quantization_for_source_paths( + "decoder", + self.HF_COMPONENT_SOURCES["decoder"], + ), + component_quantization=None, + ) + self.encoder = T5Encoder(encoder_config) + self.decoder = T5Decoder(decoder_config) def preprocess_weights( self, state_dict: dict[str, torch.Tensor] 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..36a277269 --- /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_attribute_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