Repository navigation
Fold static weight layout chains after weight loading - #788
Closed
titaiwangms wants to merge 1 commit into
Closed
titaiwangms wants to merge 1 commit into
titaiwangms wants to merge 1 commit into
Conversation
Add a generic lazy Reshape/Transpose initializer pass for higher-rank weight layouts that the existing two-dimensional transpose pass misses. Preserve dtype bits, output names, and tied-table sharing while accepting static Constant shape vectors and cleaning dead inputs in their own graph. Integrate the pass into the existing main post-weight pipeline without including the Gemma4 VLM changes. Cover values, laziness, invalid shapes, subgraph ownership, and package persistence, including BF16 special bits. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9a2a5b82-7cd3-4a76-bb94-0998527ea723 Signed-off-by: titaiwang <titaiwang@microsoft.com>
Performance Comparison
|
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Standalone folding leaves dead parent-owned initializers registered when processing subgraph captures, allowing duplicate weight serialization.
1 open issue
What changed in this PR
Adds model-independent folding of static weight-layout chains after weight loading, extending existing optimizations to higher-rank tensors without including the Gemma4-specific work from #787.
Changes:
- Lazily folds initializer-backed Reshape/Transpose chains while preserving dtype and values.
- Integrates the pass into the existing post-weight-loading pipeline.
- Adds coverage for layout values, laziness, sharing, invalid inputs, subgraphs, and package persistence.
| File | Description |
|---|---|
src/mobius/_passes/_fold_weight_layout.py |
Implements lazy static weight-layout folding and input cleanup. |
src/mobius/_passes/_fold_weight_layout_test.py |
Adds correctness, preservation, and persistence tests. |
src/mobius/_passes/__init__.py |
Exports and documents the new pass. |
src/mobius/_optimizations.py |
Adds the pass before post-weight dead-code removal. |
Review effort: Balanced
justinchuby
approved these changes
Oct 23, 2026
3 of 9 tasks
Contributor
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
Add a model-independent
FoldWeightLayoutPassto the post-weight-loading pipeline. It lazily folds initializer-backedReshape/Transposechains, including higher-rank packed expert weights, scales, and zero points that the existing 2-D transpose pass does not handle.Constant(value/value_ints)shape vectors, including ONNX zero/inferred-dimension semantics.Independence from #787
This PR targets
maindirectly and does not contain or require the Gemma4 VLM commit from #787. The dtype helper,ModelPackage.save/load, and post-weight pipeline already exist onmain.The existing
mainpass order is preserved: lift constants, fold concatenations, fold transposes, run the new layout pass, then remove dead nodes. The VLM-specific pipeline reordering, runtime gates, and weight-loading changes are not included.Validation
Validated in an isolated worktree based on
main(d69efe03):git diff --checkpass._optimizations.pyalso reports an existingrules_or_passassignment-type error; the identical diagnostic was reproduced against the unmodifiedorigin/mainfile using MyPy's shadow-file baseline check. This PR does not change that loop or suppress the error.Local GPU experiment
Separately applied the identical pass to a frozen Gemma4-26B-A4B-IT INT4 export with 30 GQA nodes (
do_rotary=1). This is a frozen-artifact A/B experiment using the existing custom native runtime, not a freshmainmodel export or a claim that the Gemma4-specific runtime work is included here.Three measured requests per context, after two warmups:
Observed model-load time changed from approximately 43-44 seconds to 4.6-5.0 seconds. These are single load observations, not a controlled cold-I/O benchmark. Decode/TTFT are broadly unchanged; the three trials do not establish statistical significance. VRAM is a 20 ms sampled lower bound and does not consistently decrease.
This checks layout-only equivalence. Previously failed strict HuggingFace numerical-parity gates are unchanged and unwaived.