fix: keep requires_grad when re-wrapping a lifted parameter - #4480
Open
shoumikhin wants to merge 1 commit into
Open
fix: keep requires_grad when re-wrapping a lifted parameter#4480shoumikhin wants to merge 1 commit into
shoumikhin wants to merge 1 commit into
Conversation
shoumikhin
force-pushed
the
quantized-param-constant
branch
from
August 12, 2026 05:07
9be3c2d to
74d1a2b
Compare
shoumikhin
force-pushed
the
quantized-param-constant
branch
from
August 12, 2026 06:18
74d1a2b to
a78a49d
Compare
Lifting a graph attribute into a placeholder re-wraps each named parameter:
state_dict[name] = torch.nn.Parameter(state_dict[name])
gm.state_dict() returns plain tensors, so the original requires_grad flag is
already gone by this point, and Parameter defaults it to True. Two problems follow.
An integer parameter cannot have requires_grad=True at all, so export fails:
RuntimeError: Only Tensors of floating point and complex dtype can require
gradients
A frozen float parameter silently comes back trainable, which builds an autograd
graph in a program meant only for inference.
The loop already iterates the module parameters, so carry the flag over from the
one it just matched instead of discarding it. Nothing else about the
classification changes: the weight stays a PARAMETER, stays in the state dict as
a Parameter, and stays out of the constants mapping, so loading a checkpoint into
the exported module keeps working. This matches torch.export, which takes the
input kind from module registration rather than from dtype.
Extends the existing lift() parameter test to assert the kind, the container type
and the requires_grad flag, and covers both flag values so an implementation that
hardcodes either one fails. Verified against the current code and against three
wrong implementations:
correct fix 4 passed
hardcode requires_grad=False 1 failed
hardcode requires_grad=True 3 failed
read the flag off the state dict 1 failed
current code, no fix 3 failed
Full file: 20 passed.
shoumikhin
force-pushed
the
quantized-param-constant
branch
from
August 12, 2026 23:47
a78a49d to
e0e0d24
Compare
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.
The problem
Legacy export fails when a registered parameter has an integer dtype. It breaks while the
exporter turns the graph attributes into placeholders.
Each named parameter is re-wrapped there:
gm.state_dict()hands back plain tensors, so the originalrequires_gradflag isalready gone by this point, and
Parameterputs it back asTrueby default. An integertensor cannot require gradients, so the export stops:
This is easy to hit without any quantizer involved: constant folding registers folded
constants as parameters, so an integer constant expression reaches the same line.
There is a second, quieter problem. A float parameter the model had frozen comes back
trainable, which builds an autograd graph inside a program meant only for inference.
The fix
The loop already has the module's parameter in hand, so carry its flag across instead of
throwing it away:
Nothing else about the classification changes. The weight stays a
PARAMETER, stays inthe state dict as a
Parameter, and stays out of the constants mapping, so loading acheckpoint into the exported module keeps working. This matches
torch.export, whichtakes the input kind from module registration rather than from dtype.
Testing
Extended the existing
lift()parameter test to assert the input kind, the container typeand the
requires_gradflag, not just the dtype and stride. It covers both flag values, soan implementation that hardcodes either one fails.
Checked against the current code and against three wrong implementations:
That last mutant matters:
state_dict()detaches, so reading the flag from it alwaysyields
Falseand looks correct until a trainable parameter is exported.Full file: 20 passed.
These tests live in the
executorchsuite, which only the nightly lane selects, so they donot run on this PR by default.