Always lift buffers when importing frozen programs#4658
Open
j2kun wants to merge 3 commits into
llvm:mainllvm/torch-mlir:mainfrom
j2kun:lift-and-freeze-state-dict-buffersj2kun/torch-mlir:lift-and-freeze-state-dict-buffersCopy head branch name to clipboard
Open
Always lift buffers when importing frozen programs#4658j2kun wants to merge 3 commits intollvm:mainllvm/torch-mlir:mainfrom j2kun:lift-and-freeze-state-dict-buffersj2kun/torch-mlir:lift-and-freeze-state-dict-buffersCopy head branch name to clipboard
j2kun wants to merge 3 commits into
llvm:mainllvm/torch-mlir:mainfrom
j2kun:lift-and-freeze-state-dict-buffersj2kun/torch-mlir:lift-and-freeze-state-dict-buffersCopy head branch name to clipboard
Conversation
sahas3
approved these changes
Jul 20, 2026
sahas3
left a comment
Member
There was a problem hiding this comment.
LG.
The PyTorch version failure was just fixed in main. Can you rebase and rerun the CIs? Thanks!
Co-authored-by: Sayan Saha <sayans@mathworks.com>
j2kun
force-pushed
the
lift-and-freeze-state-dict-buffers
branch
from
July 23, 2026 21:59
48b3f3d to
dff6baa
Compare
j2kun
enabled auto-merge (squash)
July 23, 2026 22:01
Contributor
Author
|
The failure is vs torch-nightly, not sure if that's blocking submission, as it doesn't seem related to the changes in this PR. |
Member
|
All CIs need to be passing for merging the change. It looks like some test in |
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
2 participants
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.
Reproduction: https://gist.github.com/j2kun/c2539e1fd9a7de0d10e9513cff0e989c
In fx_importer.py, import_frozen_program only lifts buffers from
state_dictto inlined constants in the else block ofif hasattr(prog, "constants"):It should always lift these buffers since the program is frozen. I ran into this because I was exporting a torch model that had batch norms in it, and the batch norm parameters were left as (many many) function arguments when I wanted to freeze them from the values they got during training.
(Followed up from Discord)