`torch.compiler.disable()` on module hooks will disable `module.compile()` #142358

gau-nernst · 2024-12-09T09:00:53Z

🐛 Describe the bug

torch.compiler.disable() on module hooks will disable compile on the whole module when module.compile() is used (or via torch.compile(some_function)(model, some_inputs))

import torch
from torch import nn


d_model = 16
model = nn.Sequential(
    nn.Linear(d_model, d_model),
    nn.GELU(),
    nn.Linear(d_model, d_model),
)
model.cuda()

@torch.compiler.disable()
def hook(module, input):
    pass

model.register_forward_pre_hook(hook)
model.compile()  # this doesn't work
# model.forward = torch.compile(model.forward)  # this works

x = torch.randn(4, d_model).cuda()
model(x)

To check if the model is compiled, I use TORCH_LOGS="output_code". For model.compile(), there is no code output. For model.forward = torch.compile(model.forward), there is compiled triton code as expected.

Although model.forward = torch.compile(model.forward) is a valid workaround, it prevents me from doing full-compile a model in a function e.g. compute loss with signature compute_loss(model, inputs, labels). Again, for this case, a workaround is to isolate the loss computation logic and compile it separately, but it gets messy. Ideally torch.compile should only ignore the hooks, not the .forward() logic, correctly.

Error logs

No response

Versions

torch==2.6.0.dev20241208+cu126

cc @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @amjames

The text was updated successfully, but these errors were encountered:

vadimkantorov · 2024-12-09T14:02:04Z

I wonder if it's related to global-ness of compiler options... (#124505)

ezyang · 2024-12-09T14:55:45Z

If someone wants to dig in, TORCH_LOGS=+dynamo should give detailed information about why we are skipping frames and should reveal the problem. tlparse doesn't do so good with telling us about disable interactions.

gau-nernst · 2024-12-10T00:59:33Z

@ezyang I created a gist here with TORCH_LOGS="+dynamo" output https://gist.github.com/gau-nernst/0d36db3a3331b3f83548e5b46370ae17

ezyang · 2024-12-11T19:32:22Z

< AD5E tr class="d-block">

@anijain2305 this feels related to the fact that we don't trace __call__ and we wait until we hit forward.

@gau-nernst I think a hack that will workaround the problem for you is to make a function which wraps the Sequential call.

gau-nernst · 2024-12-12T00:59:08Z

@ezyang Thank you for getting back. I have tried but this doesn't work

def f(module, input):
    return module(input)

x = torch.randn(4, d_model).cuda()
torch.compile(f)(model, x)

This also doesn't work, which is a typical transformer architecture, and I want to install hooks on each transformer layer

class Model(nn.Module):
    def __init__(self, d_model: int):
        super().__init__()
        self.embed = nn.Embedding(10, d_model)
        self.layers = nn.ModuleList([nn.Linear(d_model, d_model) for _ in range(2)])
        self.output = nn.Linear(d_model, 10)

    def forward(self, x):
        x = self.embed(x)
        for layer in self.layers:
            x = layer(x)
        x = self.output(x)
        return x


model = Model(d_model).cuda()
for layer in model.layers:
    layer.register_forward_pre_hook(hook)

model.compile()
x = torch.randint(10, size=(4, 16)).cuda()
model(x)

Even the embedding and output layer are not compiled. Though the "+dynamo" output is different (it's easy to reproduce, so I think I don't need to paste it here).

gau-nernst added the oncall: pt2 label Dec 9, 2024

ezyang added module: dynamo module: nn Related to torch.nn triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module labels Dec 9, 2024

ezyang added the triage review label Dec 11, 2024

IvanKobzarev removed the triage review label Dec 17, 2024

anijain2305 self-assigned this Jan 17, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

`torch.compiler.disable()` on module hooks will disable `module.compile()` #142358

`torch.compiler.disable()` on module hooks will disable `module.compile()` #142358

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

torch.compiler.disable() on module hooks will disable module.compile() #142358

Uh oh!

🐛 Describe the bug

Error logs

Versions

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

`torch.compiler.disable()` on module hooks will disable `module.compile()` #142358