Migrate composite model weight adapters - #670
Draft
xiaoyu-work wants to merge 2 commits into
Draft
Conversation
Performance Comparison
|
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
xiaoyu-work
force-pushed
the
feat/per-component-quantization-loader
branch
from
August 27, 2026 21:08
e409b2d to
a8a5c5b
Compare
xiaoyu-work
force-pushed
the
refactor/component-weight-adapters
branch
from
August 27, 2026 21:10
ea293cb to
18bdcfd
Compare
xiaoyu-work
force-pushed
the
feat/per-component-quantization-loader
branch
from
August 27, 2026 21:22
a8a5c5b to
7562458
Compare
xiaoyu-work
force-pushed
the
refactor/component-weight-adapters
branch
from
August 27, 2026 21:28
18bdcfd to
b4e4bb8
Compare
xiaoyu-work
force-pushed
the
feat/per-component-quantization-loader
branch
from
August 27, 2026 21:30
7562458 to
e5ce365
Compare
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 <[email protected]>
Update binding validation and adapter fixtures for ComponentDescriptor.module_attribute_path. Signed-off-by: Xiaoyu Zhang <[email protected]>
xiaoyu-work
force-pushed
the
refactor/component-weight-adapters
branch
from
August 27, 2026 21:31
b4e4bb8 to
f998a0b
Compare
Contributor
There was a problem hiding this comment.
🔵 Needs a closer look
It changes core weight-loading/quantization behavior across multiple model families and build entrypoints, so it warrants final human review despite strong test coverage.
Pull request overview
This PR introduces an explicit model-weight adapter boundary in the checkpoint loading pipeline (while preserving legacy preprocess_weights), updates several model implementations to respect per-component quantization plans, and adds a binding gate that fails builds/publication when affine-quantized ops still have unbound parameters after weight application.
Changes:
- Add
adapt_model_weights()+ModelWeightAdapter/WeightAdapterContextas the new narrow model-specific adapter boundary with legacy fallback. - Make Qwen3.5/QMoE, Gemma4, and T5 preprocessing/component-quantization interactions component-plan aware (including routing and aliasing).
- Add
validate_quantized_component_bindings()and invoke it in build paths to ensure MatMulNBits/GatherBlockQuantized parameters are actually bound.
File summaries
| File | Description |
|---|---|
| src/mobius/weights/_adapters.py | Introduces the explicit model-weight adapter boundary with legacy preprocess_weights fallback. |
| src/mobius/weights/_adapters_test.py | Adds tests covering legacy hook compatibility and explicit adapter precedence. |
| src/mobius/weights/init.py | Exposes adapter API from the mobius.weights public surface. |
| src/mobius/models/t5.py | Adjusts tied-weight handling when component_quantization is active to avoid double preprocessing. |
| src/mobius/models/qwen35.py | Routes decoder quantization decisions through the decoder component plan and avoids double preprocessing in component mode. |
| src/mobius/models/qwen35_test.py | Adds regression test ensuring plan-only quantization keeps QMoE graph/weights aligned. |
| src/mobius/models/gemma4.py | Makes decoder/vision/audio/embedding construction and routing component-quantization aware; adds HF alias metadata. |
| src/mobius/models/gemma4_test.py | Adds component-plan layout, aliasing, routing, and embedding-sidecar routing regressions. |
| src/mobius/integrations/transformers/_builder.py | Switches weight preprocessing to adapt_model_weights() and enforces post-load binding validation. |
| src/mobius/components/_moe.py | Uses decoder component quantization when selecting QMoE quantization compatibility. |
| src/mobius/_component_quantization.py | Extends module→HF-name resolution (incl. .linear alias), allows specialized quantized modules when layouts match, and adds binding validator. |
| src/mobius/_component_quantization_test.py | Adds test ensuring the binding validator rejects unfilled affine-quantized parameters. |
| src/mobius/main.py | Applies adapter boundary in CLI load path and enforces binding validation post-weight-load. |
| CHANGELOG.md | Documents the adapter migration and component-sidecar routing/binding behavior. |
Review details
- Files reviewed: 14/14 changed files
- Comments generated: 1
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+59
to
+61
| preprocess = getattr(module, "preprocess_weights", None) | ||
| if preprocess is None: | ||
| return dict(state_dict) |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
preprocess_weightscompatibilityStack
Validation
.linearalias, per-component layouts, sidecar routing, and quantized embedding regressions