Add typed checkpoint weight records - #666
Closed
xiaoyu-work wants to merge 8 commits into
Closed
Conversation
Performance Comparison
|
xiaoyu-work
force-pushed
the
refactor/typed-weight-pipeline
branch
from
August 27, 2026 02:21
b2b820e to
fabb835
Compare
xiaoyu-work
force-pushed
the
refactor/typed-weight-pipeline
branch
from
August 27, 2026 02:41
fabb835 to
b1ea32f
Compare
xiaoyu-work
force-pushed
the
refactor/typed-weight-pipeline
branch
from
August 27, 2026 21:06
b1ea32f to
e02875f
Compare
xiaoyu-work
force-pushed
the
refactor/component-manifest
branch
from
August 27, 2026 21:13
b549476 to
501c342
Compare
xiaoyu-work
force-pushed
the
refactor/typed-weight-pipeline
branch
from
August 27, 2026 21:13
e02875f to
ea9c144
Compare
Resolve task roles, module paths, and HuggingFace source ownership into one immutable manifest. Switch inspection and optimization-role lookup to the shared metadata without changing graph or weight behavior. Signed-off-by: Xiaoyu Zhang <xiaoyuzhang@microsoft.com>
Derive candidate HuggingFace module names from component-local paths so later loader stages can apply exact and regex quantization exclusions consistently. Signed-off-by: Xiaoyu Zhang <xiaoyuzhang@microsoft.com>
Only derive HuggingFace module candidates when a declared source root shares an anchor with the component-local path, preventing unrelated roots such as lm_head from matching every decoder module. Signed-off-by: Xiaoyu Zhang <xiaoyuzhang@microsoft.com>
Allow model declarations to map component-local module prefixes to HuggingFace source prefixes when structural anchor inference is insufficient. Signed-off-by: Xiaoyu Zhang <xiaoyuzhang@microsoft.com>
Rename module_path to module_attribute_path and document task-defined optimization roles so callers cannot confuse Python module ownership with package or checkpoint prefixes. Signed-off-by: Xiaoyu Zhang <xiaoyuzhang@microsoft.com>
xiaoyu-work
force-pushed
the
refactor/component-manifest
branch
from
August 27, 2026 21:30
501c342 to
1f47daa
Compare
Group affine packed sidecars into component-owned logical records and introduce a quantization codec registry over the existing Olive, GPTQ, and AWQ normalization helpers. The active loader remains unchanged. Signed-off-by: Xiaoyu Zhang <xiaoyuzhang@microsoft.com>
Update typed weight fixtures for the clarified ComponentDescriptor API. Signed-off-by: Xiaoyu Zhang <xiaoyuzhang@microsoft.com>
xiaoyu-work
force-pushed
the
refactor/typed-weight-pipeline
branch
from
August 27, 2026 21:30
ea9c144 to
15c0dfc
Compare
Member
Author
|
@copilot resolve the merge conflicts in this pull request |
…-pipeline # Conflicts: # src/mobius/_builder.py Co-authored-by: xiaoyu-work <85524621+xiaoyu-work@users.noreply.github.com>
Contributor
Resolved in a471d34 (merge commit 010ae62): merged origin/main, conflict was in |
|
|
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
| config: QuantizationConfig, | ||
| ) -> WeightBundle: | ||
| """Group checkpoint sidecars into typed logical records.""" | ||
| ... |
| config: QuantizationConfig, | ||
| ) -> dict[str, torch.Tensor]: | ||
| """Convert one packed record to Mobius's canonical parameter layout.""" | ||
| ... |
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
The new codec and record abstractions have missing consistency checks that can allow mismatched quant_method normalization and inconsistent packed-sidecar state.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
This PR introduces a typed “logical weight” layer in mobius.weights to represent float tensors and packed affine-quantized sidecars as first-class records, plus a codec registry keyed by quant_method to group and normalize existing Olive/GPTQ/AWQ checkpoints without changing the active loader.
Changes:
- Add typed record abstractions:
FloatWeight,PackedWeight,WeightRecord, and component-scopedWeightBundle. - Add
QuantizationCodec+QuantizationCodecRegistry, with a legacy codec that groups packed sidecars and normalizes them through the existingpreprocess_quantized_weightshelpers. - Add unit tests covering record/bundle invariants and codec grouping/registry behavior.
File summaries
| File | Description |
|---|---|
| src/mobius/weights/_records.py | Adds typed logical weight record dataclasses and an immutable bundle wrapper. |
| src/mobius/weights/_records_test.py | Unit tests for WeightBundle routing and source key tracking. |
| src/mobius/weights/_codecs.py | Adds codec protocol + registry and legacy Olive/GPTQ/AWQ grouping/normalization implementation. |
| src/mobius/weights/_codecs_test.py | Unit tests for grouping behavior, error cases, normalization, and registry semantics. |
| src/mobius/weights/init.py | Exposes the new typed weight and codec APIs via mobius.weights. |
Review details
- Files reviewed: 5/5 changed files
- Comments generated: 2
- 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
+189
to
+194
| return preprocess_quantized_weights( | ||
| record.storage.as_state_dict(), | ||
| config, | ||
| tie_embeddings=False, | ||
| qmoe_target_path=None, | ||
| ) |
Comment on lines
+36
to
+43
| qweight: torch.Tensor | ||
| scales: torch.Tensor | ||
| zero_points: torch.Tensor | None | ||
| qweight_key: str | ||
| scales_key: str | ||
| zero_points_key: str | None | ||
| method: str | ||
|
|
Member
Author
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
WeightRecord,PackedWeight, and component-ownedWeightBundleabstractionsquant_methodStack
Validation