Merged main

[micwill755]

Description

Dynamo constant folding can replace input-independent factories (e.g. torch.zeros) with persistent _frozen_param* attributes. When those attributes are also graph outputs and eager code after a graph break mutates them in-place, later invocations observe leftover state instead of fresh tensors.

This change clones folded _frozen_param* values at the graph-output boundary so downstream in-place mutation cannot modify compiler-owned state, preserving eager semantics across calls.

Fixes #4466

Type of change

• Bug fix (non-breaking change which fixes an issue)

Checklist:

• My code follows the style guidelines of this project (You can use the linters)
• I have performed a self-review of my own code
• I have commented my code, particularly in hard-to-understand areas and hacks
• I have made corresponding changes to the documentation
• I have added tests to verify my fix or my feature
• New and existing unit tests pass locally with my changes
• I have added the relevant labels to my PR so that relevant reviewers are notified

[lanluo-nvidia]

What codex review says:

• P1 — Preserve aliases for repeated folded outputs. _clone_folded_constants_at_outputs creates a new clone for every output occurrence. If a graph returns the same folded tensor twice (return x, x), the PR changes eager semantics: the two
  results no longer share storage. Cache one clone per get_attr node and reuse it for repeated outputs; add a regression test that mutates one returned value and observes the other.

[narendasan]

@micwill755 can you explain a bit more about this?

Dynamo constant folding can replace input-independent factories (e.g. torch.zeros) with persistent _frozen_param* attributes. When those attributes are also graph outputs and eager code after a graph break mutates them in-place, later invocations observe leftover state instead of fresh tensors.

[micwill755]

@micwill755 can you explain a bit more about this?

Dynamo constant folding can replace input-independent factories (e.g. torch.zeros) with persistent _frozen_param* attributes. When those attributes are also graph outputs and eager code after a graph break mutates them in-place, later invocations observe leftover state instead of fresh tensors.

Constant folding turns torch functions like torch.zeros(...) into a persistent module attribute i.e. _frozen_param0 which is reused across calls. If that attribute is also returned from the graph, eager code after a graph break can mutate it in-place, e.g. x += 1, which edits the stored constant. This means the next call then sees that leftover state instead of a fresh tensor of zeros, unlike normal eager PyTorch, which allocates new storage every time.

[narendasan]

What is the case where a constant is returned as output?

[micwill755]

What is the case where a constant is returned as output?

If an input-independent factory op e.g. torch.zeros is folded into a frozen_param and that value is also a graph return, then a constant is returned as an output. A graph break is a common reason this happens because the the factory lives in the compiled subgraph, and eager code after the break can mutate that returned constant in-place, impacting later calls.

[cehongwang]

The motivation makes sense. Instead of doing the clone, can we disable the constant folding for the frozen parameter that goes to the output? Also, in AoT compilation, when we are doing constant folding, the graph has not been split yet (the partition happens after the constant folding). After the partition, there could be new subgraphs that add output to the intermediate subgraphs. So we may need to adjust the frozen_parameter after the partitioning.

[narendasan]

@micwill755 I thought we talked about the distinction between module state and function state

class Mod(nn.Module):
    def __init__(self):
        self.weight = torch.zeros(...)

    def forward(self, x):
         

class Mod(nn.Module):
    def __init__(self):


    def forward(self, x):
        self.weight = torch.zeros(...)

these two modules would have different semantics wrt mutation between calls right? Is this handled if we just pattern match on get_attr / frozen?

[narendasan]

Also we talked about if we should clone at construction or clone on return right?