Unnecessary compilation fails to optimize simple code

[youkaichao]

🐛 Describe the bug

A minimal reproducible example:

import torch

class Layer(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.weight = torch.randn((16,))
        self.variance_epsilon = 1e-5

    @torch.compile
    def forward(self, hidden_states, residuals=None):
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        mean = hidden_states.mean(-1, keepdim=True)
        variance = (hidden_states - mean).pow(2).mean(-1, keepdim=True)
        hidden_states = (hidden_states -
                            mean) * torch.rsqrt(variance + self.variance_epsilon)
        hidden_states = self.weight.to(torch.float32) * hidden_states
        return hidden_states.to(input_dtype), residuals

layers = [Layer() for i in range(100)]
hidden_states = torch.randn((32, 16, 16))

for iteration in range(2):
    # simulate a model forward call
    for layer in layers:
        hidden_states, _ = layer(hidden_states)

For these 100 layers, torch.compile will compile the first 64 layers (which is the dynamo size limit for a code object), and the rest layers are not optimized.

However, ideally, we should only have one cache entry, that can be shared for all layers. We don't need to create different cache entry depending on id(self).

cc @ezyang @msaroufim @bdhirsh @anijain2305 @chauhang @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @jansel

Error logs

No response

Minified repro

No response

Versions

pytorch 2.3.0+cu121

[bdhirsh]

The problem is that dynamo burns the parameters/buffers into the graph that we compile for each layer, forcing them to get specialized.

@anijain2305 recently added an (experimental?) config to avoid that burning in, that you can try by running with TORCHDYNAMO_INLINE_INBUILT_NN_MODULES=1.

I actually tried it on your repro, and I (successfully) don't see any recompiles when I turn it on.

[youkaichao]

@bdhirsh Can you explain the rationale of the specialization? I'm quite puzzled here.

[ezyang]

This is Animesh's thing, he's working on fixing it

[youkaichao]

@bdhirsh thanks for the information. I suppose that requires several months to be public, right?

[ezyang]

@anijain2305 seemed pretty close when we talked about it a week ago

[ezyang]

in particular, the flag is already available, you can opt into it and see if it works

[youkaichao]

Thanks for the answer. Setting export TORCHDYNAMO_INLINE_INBUILT_NN_MODULES=1 indeed solves this particular problem.

May I ask why this is not the default? Why do we need to set it manually?

[ezyang]

@anijain2305 is working on setting it default on. It currently uncovers a pile of latent bugs that are showing on test suite.