Skip to content

fix: keep requires_grad when re-wrapping a lifted parameter - #4480

Open
shoumikhin wants to merge 1 commit into
pytorch:mainfrom
shoumikhin:quantized-param-constant
Open

fix: keep requires_grad when re-wrapping a lifted parameter#4480
shoumikhin wants to merge 1 commit into
pytorch:mainfrom
shoumikhin:quantized-param-constant

Conversation

@shoumikhin

@shoumikhin shoumikhin commented Aug 12, 2026

Copy link
Copy Markdown
Contributor

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:

state_dict[name] = torch.nn.Parameter(state_dict[name])

gm.state_dict() hands back plain tensors, so the original requires_grad flag is
already gone by this point, and Parameter puts it back as True by default. An integer
tensor cannot require gradients, so the export stops:

RuntimeError: Only Tensors of floating point and complex dtype can require gradients

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:

for name, param in gm.named_parameters():
    if node.target == name:
        input_kind = InputKind.PARAMETER
        state_dict[name] = torch.nn.Parameter(
            state_dict[name], requires_grad=param.requires_grad
        )
        break

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.

Testing

Extended the existing lift() parameter test to assert the input kind, the container type
and the requires_grad flag, not just the dtype and stride. It covers both flag values, so
an implementation that hardcodes either one fails.

Checked 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

That last mutant matters: state_dict() detaches, so reading the flag from it always
yields False and looks correct until a trainable parameter is exported.

Full file: 20 passed.

These tests live in the executorch suite, which only the nightly lane selects, so they do
not run on this PR by default.

@meta-cla meta-cla Bot added the cla signed label Aug 12, 2026
@github-actions github-actions Bot added component: core Issues re: The core compiler component: api [Python] Issues re: Python API component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths labels Aug 12, 2026
@shoumikhin
shoumikhin force-pushed the quantized-param-constant branch from 9be3c2d to 74d1a2b Compare August 12, 2026 05:07
@github-actions github-actions Bot added the component: tests Issues re: Tests label Aug 12, 2026
@shoumikhin shoumikhin changed the title fix: classify an integer weight as a constant, not a parameter fix: keep requires_grad when re-wrapping a lifted parameter Aug 12, 2026
@shoumikhin
shoumikhin force-pushed the quantized-param-constant branch from 74d1a2b to a78a49d Compare August 12, 2026 06:18
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
shoumikhin force-pushed the quantized-param-constant branch from a78a49d to e0e0d24 Compare August 12, 2026 23:47
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cla signed component: api [Python] Issues re: Python API component: core Issues re: The core compiler component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths component: tests Issues re: Tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant