Load quantized checkpoint weights per component - #733
Conversation
Consolidate the typed weight records, per-component loader, and composite model adapters from #666, #669, and #670 into one change based on current main. The design and canonical component manifests are already merged through #664 and #665. Preserve independent component layouts, source-name exclusions, Olive/GPTQ/AWQ sidecar normalization, Gemma4/Qwen3.5/T5 model adaptations, and quantized parameter binding validation. Clean up import, type-annotation, and formatting issues left by the earlier stacked-branch conflict resolutions. Signed-off-by: Xiaoyu Zhang <xiaoyuzhang@microsoft.com>
Performance Comparison
|
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
There was a problem hiding this comment.
🟡 Changes recommended
There are a few correctness/maintenance issues in the updated component-quantization plumbing (including an optional-arg path that will fail at runtime and redundant/accidental code) that should be fixed before approval.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
This PR upgrades Mobius’s quantized-checkpoint loading pipeline to be component-aware, introducing typed logical weight records and a codec registry to normalize existing Olive/GPTQ/AWQ packed sidecars per component (without quantizing float checkpoints), while preserving model-specific weight adaptation behavior and validating that quantized graph parameters are actually bound before publishing.
Changes:
- Add
mobius.weightstyped abstractions (records/bundles), a quantization codec registry, and an explicit model-weight adapter boundary. - Switch Transformers + CLI build paths to route weights through the adapter + per-component normalization, then enforce a quantized-parameter binding gate.
- Update per-component quantization/config plumbing and extend Gemma4/Qwen3.5/QMoE/T5 handling + tests to keep graph layouts and loaded packed weights aligned.
File summaries
| File | Description |
|---|---|
| tests/build_graph/_support.py | Adds helper to assign per-component quantization layouts in graph tests. |
| src/mobius/weights/_records.py | Introduces typed logical weight records (FloatWeight, PackedWeight, WeightRecord, WeightBundle). |
| src/mobius/weights/_records_test.py | Unit tests for typed weight records and bundle invariants. |
| src/mobius/weights/_codecs.py | Adds codec interface + registry and a legacy affine codec wrapper (Olive/GPTQ/AWQ). |
| src/mobius/weights/_codecs_test.py | Tests for codec grouping/validation/normalization behavior. |
| src/mobius/weights/_adapters.py | Adds explicit model adapter boundary with legacy preprocess_weights fallback. |
| src/mobius/weights/_adapters_test.py | Tests adapter precedence and legacy hook compatibility. |
| src/mobius/weights/init.py | Exposes new typed weight/codec/adapter APIs from the mobius.weights package. |
| src/mobius/models/t5.py | Applies per-component quantization plans to encoder vs decoder construction. |
| src/mobius/models/qwen35.py | Routes decoder quantization through component plans and aligns VL decoder source paths. |
| src/mobius/models/qwen35_test.py | Adds regression test for plan-only quantization alignment with QMoE packing. |
| src/mobius/models/gemma4.py | Extends component quantization resolution and adds HF module alias metadata for routing/exclusions. |
| src/mobius/models/gemma4_test.py | Adds broader component-plan tests: regex exclusions, layout-per-component, routing, embedding sidecars. |
| src/mobius/integrations/transformers/_builder.py | Switches loading to use component manifests + adapters + per-component normalization + binding validation. |
| src/mobius/integrations/transformers/_builder_test.py | Adjusts test config to reflect per-component layouts and stripping behavior. |
| src/mobius/components/_moe.py | Ensures QMoE quantization compatibility uses decoder component quantization when present. |
| src/mobius/_configs/_quantization.py | Reworks module-plan rules (modules_to_not_convert, overrides) and regex validation. |
| src/mobius/_configs/_base.py | Clarifies component-quantization extraction semantics and config comments. |
| src/mobius/_configs_test.py | Updates tests for new module-rule matching and component quantization parsing behavior. |
| src/mobius/_component_quantization.py | Major refactor: manifest-driven component configuration, packed-sidecar routing/normalization, binding validation. |
| src/mobius/_component_quantization_test.py | Updates/expands tests for authoritative component plans, routing, and binding validation. |
| src/mobius/_component_manifest.py | Allows manifest resolution without requiring model_type to be non-null. |
| src/mobius/_builder.py | Threads a resolved ComponentManifest through build to keep graph+weight routing aligned. |
| src/mobius/main.py | Updates CLI build/load path to use adapter + component normalization + binding validation. |
| CHANGELOG.md | Updates Unreleased notes to describe per-component quantized checkpoint loading behavior. |
Review details
- Files reviewed: 25/25 changed files
- Comments generated: 4
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Bring all seven pending source changes from the old per-component quantization stack into the unified PR instead of leaving real feature work on a separate branch. Keep sparse per-layer overrides from replacing default head dimensions and KV-head counts. Preserve a layered per-layer-input ABI between Gemma4 embedding and decoder exports for OpenVINO while retaining the flattened layout for other execution providers. Include the accompanying regressions and document the component contract. Signed-off-by: Xiaoyu Zhang <xiaoyuzhang@microsoft.com>
Carry forward the seq2seq, vision-language, and speech component-layout matrix tests from the old integration branch. These tests were lost when the monolithic graph test file was split during earlier stacked-branch conflict resolutions, leaving their shared configuration helper unused. Reuse the shared helper in the current split test modules so independent component bit widths and group sizes remain covered by the unified PR. Signed-off-by: Xiaoyu Zhang <xiaoyuzhang@microsoft.com>
Resolve graph layout and checkpoint normalization with the same per-module policy, including HuggingFace source aliases, model float exclusions and output heads. Preserve raw sidecars until the effective layout is known, and distinguish Gather embedding tables from MatMul weight layouts. Keep legacy absent-component configuration working, defer Gemma4 component normalization, and restore exact T5 source ownership without applying decoder stack rules to a separately owned LM head. Add packed-value loading regressions for mixed components, overrides, symmetry, and protected output heads. Signed-off-by: Xiaoyu Zhang <xiaoyuzhang@microsoft.com>
Match literal quantization rules on module-path boundaries and require an explicit task for compatibility normalization. Remove redundant Protocol placeholders and simplify Gemma4 source aliases. Use a shared scaled quantized embedding for Gemma3 component plans while preserving Gemma4 behavior. Normalize materialized split tied tables against supported targets, retain unsafe-tie rejection, and cover actual embedding execution as well as effective head-layout mismatches. Correct the LLaVA vision fixture namespace and update the stale GGUF runtime admission expectation without weakening tokenizer identity or numerical parity checks. The focused suite and pinned SmolLM2 GGUF integration pass. Signed-off-by: Xiaoyu Zhang <xiaoyuzhang@microsoft.com>
There was a problem hiding this comment.
Note
Copilot was unable to run its full agentic suite in this review.
Pull request overview
Copilot reviewed 41 out of 41 changed files in this pull request and generated 5 comments.
Suppressed comments (1)
tests/build_graph/_support.py:1
- This test helper assigns layouts based on
enumerate(task.model_roles.items()), so the chosen (bits, group_size) depends on the insertion order ofmodel_roles. Ifmodel_rolesdiffers across tasks or is edited, tests can become brittle/non-obvious.\n\nSuggestion (optional): iterate in a stable order (e.g.,for index, component in enumerate(sorted(task.model_roles))) or explicitly define the component order you want to exercise, so layout assignment is deterministic and resilient to unrelated dictionary edits.
Resolve dynamic source hooks from an explicit or HuggingFace config model type, and skip them when neither exists instead of passing an empty identifier. Keep static sources authoritative for classes without dynamic hooks. Remove the unused model-type argument from config-based alias hooks so hand-built T5 configs retain their layer mappings. Cover missing identifiers, explicit type precedence, and static/config-based metadata preservation, and document the hook contract. Signed-off-by: Xiaoyu Zhang <xiaoyuzhang@microsoft.com>
There was a problem hiding this comment.
🟡 Changes recommended
Projection overrides, explicit float MoE gates, binding coverage, and the documented adapter contract remain inconsistent.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
Suppressed comments (1)
Previously missed (1) — in code that hasn't changed since the last review.
src/mobius/weights/_adapters.py:41
- The public adapter protocol does not match the component-owned typed pipeline documented in
docs/design/per-component-quantization-loading.md:373-402. It receives the entire raw state dict before routing/grouping instead of aComponentDescriptorandWeightBundle, so an implementation cannot enforce single-component ownership and can still consume or normalize unrelated sidecars. Either invoke adapters per routed bundle after grouping, or explicitly revise the design and PR API contract as a transitional raw-state adapter.
- Files reviewed: 41/41 changed files
- Comments generated: 3
- Review effort level: Balanced
titaiwangms
left a comment
titaiwangms
left a comment
There was a problem hiding this comment.
Full-team review of ad666a825b9becf297ae350b5a98f05dc2ae831a (readability, correctness, adversarial, semantic-invariant, and integration passes).
The typed records/codec boundary and final binding check are strong improvements, but I found four correctness issues that should be addressed before merge:
-
Every plain
Embeddingin a quantized component is now treated as a token table.
src/mobius/_component_quantization.py:346-366rewrites any ordinaryEmbeddingwhenquantize_embeddings=True. The previous implementation restricted this to known token-table names, and this PR removes the test that ensured positional embeddings remained float. Components such as speech decoders can own both token and learned positional embeddings; a positional table may therefore becomeGatherBlockQuantizedeven though the checkpoint contains only its float weight, leaving packed inputs unbound. Please make this decision manifest/source-name driven (or otherwise restore a token-table-only gate) and add a non-token embedding regression test. -
Qwen3.5 component adapters still normalize ordinary decoder sidecars before per-projection overrides are resolved.
src/mobius/models/qwen35.py:582-588,:702-722, and:988-1005callpreprocess_quantized_weightswith a decoder-wide configuration.QuantizationConfig.for_source_paths()(src/mobius/_configs/_quantization.py:256-266) resolves the source roots and strips the remaining module rules. A leaf override such asmodel.language_model.layers.0.self_attn.q_projcan therefore build an overridden graph through the generic per-module rewrite while the adapter has already packed its weight with the component default. The shared loader sees canonical keys and cannot correct the layout. Please keep non-expert sidecars raw in component mode, as Gemma4 does, and limit the adapter to rename/tie/QMoE-specific transformations. -
The graph and weight paths disagree on when legacy global module rules are active.
configure_component_quantization()gates them onmanifest.names == ("model",)(src/mobius/_component_quantization.py:428-434), whilenormalize_component_quantized_weights()gates them onlen(component_names) == 1(:610-616). Single-component tasks nameddecoder,vision_encoder, etc. can consequently normalize with exclusions/overrides that were never applied to graph construction. Please use one shared predicate based on the same manifest/package ownership contract. -
Binding validation skips that newly supported legacy module-rule path.
validate_quantized_component_bindings()returns whenevercomponent_quantization is None(src/mobius/_component_quantization.py:740-746), even though configure/normalize explicitly support a rootquantizationcontaining exclusions or overrides. Missing packed weights/scales can therefore survive the loader and fail only during session creation or inference. The validator can safely scan emitted quantized ops regardless of whether the plan came from the component mapping or the root config.
Additional issues:
normalize_component_quantized_weights()skips codec grouping when there is no raw qweight (src/mobius/_component_quantization.py:628-632), and its final residual check covers only qweight keys (:728-730). An unrouted/orphan*_scales,*_qzeros,.scales, or.qzerosentry is silently accepted. Validate all packed sidecar suffixes and add a scales-only regression test.- The design contract says
QuantizationCodec.normalize(...) -> WeightRecord, while the implemented protocol returnsdict[str, torch.Tensor](docs/design/per-component-quantization-loading.md, codec section vs.src/mobius/weights/_codecs.py:174-180). Please update the design snippet so future codecs implement the actual interface.
Positive notes: complete canonical-group detection is substantially safer than suffix-only detection; zero-point validation occurs after effective override resolution; and CLI/Transformers now share the same manifest → adapter → codec → binding pipeline.
Preserve positional embeddings and architecture-specific float shared gates when applying component plans. Use the complete manifest to keep legacy module-rule activation consistent between graph construction and normalization. Defer ordinary packed sidecars until per-projection layouts are resolved, while retaining QMoE expert packing and float table tying. Reject orphan sidecars without rejecting complete canonical affine or QMoE groups, and validate bound quantized inputs for legacy root configurations. Add regression coverage for graph layouts and loaded tensor values, and clarify the implemented codec and compatibility-adapter contracts. Signed-off-by: Xiaoyu Zhang <xiaoyuzhang@microsoft.com>
There was a problem hiding this comment.
🟡 Changes recommended
GPTQ activation-order metadata is silently bypassed, and invalid embedding layouts can produce graphs that contradict the requested component plan.
Get a fresh assessment by requesting another Copilot review.
Review details
- Files reviewed: 45/45 changed files
- Comments generated: 2
- Review effort level: Balanced
| key for key in state_dict if key.endswith(("_qweight", ".qweight")) | ||
| ) | ||
| for qweight_key in qweight_keys: | ||
| logical_name, scales_key, _, zero_points_key = self._logical_name(qweight_key) |
| if isinstance(child, Embedding) and type(child).forward is Embedding.forward: | ||
| embedding_dim = int(child.weight.shape[1]) | ||
| if ( | ||
| quantization.quantize_embeddings | ||
| and any( | ||
| name.rsplit(".", 1)[-1] in _TOKEN_EMBEDDING_NAMES for name in source_names | ||
| ) | ||
| and embedding_dim % quantization.group_size == 0 | ||
| ): |
Summary
[batch, sequence, layers, projection]form between embedding and decoder components for OpenVINO; other execution providers retain the flattened layout.Loading correctness fixes
Consolidation
Replaces #666, #669, and #670 and consolidates the related development from
test/per-component-quantization-stackinto this single PR.The branch starts directly from
mainata471d3416e9c52d5508d28730962e18c8a1053d0. The design (#664) and canonical component manifests (#665) are already merged and are not separate dependencies of this PR.This includes all seven previously uncommitted Gemma4/config/EP files from the old development worktree, not just the three original PR heads. The committed old feature stack was also compared with main plus this PR: production feature areas were retained, and its missing graph-layout matrix tests were restored in the current split test modules.
Review and CI follow-up
Addressed all eight inline review threads, including duplicate/already-resolved reports, in
190eb8a5:vision_modelnamespace without relaxing target validation.The original full L4/L5 jobs hit their explicit one-hour limits after recording failures. Those failed case IDs also occur on the exact base
a471d341, but cancellation lost exception details, so they are not claimed to be infrastructure-only or fixed. No model selection, tolerance, or timeout was weakened. Isolated GPT-2 diagnostics at the original head each executed and passed exactly one selected case: L4 https://github.com/onnxruntime/mobius/actions/runs/34308638240 and L5 https://github.com/onnxruntime/mobius/actions/runs/34308638242 . These single-case successes do not certify the all-model golden jobs.Fresh full CI for the published follow-up: https://github.com/onnxruntime/mobius/actions/runs/34309655380 .
Validation
smollm2-135m-instruct-f16GGUF CLI export/reload, full-logit, and generation integration case.