pytorch - 💡(How to fix) Fix `torch.compile fails with RecursionError on nn.Module with custom __getattr__, while eager execution succeeds`

torch version: 2.12.0a0+git775500a
input shape: (1, 10)
input dtype: torch.float32
model class: MyModel

[EAGER] run succeeded.
[EAGER] output shape: (1, 10)
[EAGER] output dtype: torch.float32
[EAGER] output: tensor([[-0.1188, -0.0712, -0.0997, -0.1479,  0.1619,  0.0095,  0.0335, -0.2443,
         -0.0854,  0.2605]])

---

Traceback (most recent call last):
  File "for_test_1.py", line 36, in __getattr__
    return super().__getattr__(name)
  File "/data/zyzhao/pytorch-source/pytorch-main/torch/nn/modules/module.py", line 1967, in __getattr__
    raise AttributeError(
AttributeError: 'MyModel' object has no attribute 'custom_attributes'

During handling of the above exception, another exception occurred:

  File "for_test_1.py", line 38, in __getattr__
    if hasattr(self, "custom_attributes") and name in self.custom_attributes:
  File "for_test_1.py", line 38, in __getattr__
    if hasattr(self, "custom_attributes") and name in self.custom_attributes:
  File "for_test_1.py", line 38, in __getattr__
    if hasattr(self, "custom_attributes") and name in self.custom_attributes:
  [Previous line repeated many times]

RecursionError: maximum recursion depth exceeded while calling a Python object

---

import torch
import torch.nn as nn


class BaseModel(nn.Module):
    def __init__(self, *args, **kwargs):
        super(BaseModel, self).__init__()
        self.common_components = self.initialize_common_components(*args, **kwargs)
        self.additional_initialization(*args, **kwargs)

    def initialize_common_components(self, *args, **kwargs):
        raise NotImplementedError("Subclasses should implement this method")

    def additional_initialization(self, *args, **kwargs):
        pass

    def forward(self, *args, **kwargs):
        preprocessed_input = self.preprocess_input(*args, **kwargs)
        output = self.core_computation(preprocessed_input)
        final_output = self.post_process(output)
        return final_output

    def preprocess_input(self, *args, **kwargs):
        return args[0]

    def core_computation(self, input):
        raise NotImplementedError("Subclasses should implement this method")

    def post_process(self, output):
        return output

    def __getattr__(self, name):
        try:
            return super().__getattr__(name)
        except AttributeError:
            if hasattr(self, "custom_attributes") and name in self.custom_attributes:
                return self.custom_attributes[name]
            raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{name}'")

    def setup_loss_function(self, loss_fn=None):
        if loss_fn is None:
            self.loss_fn = nn.CrossEntropyLoss()
        else:
            self.loss_fn = loss_fn

    def setup_optimizer(self, optimizer=None, lr=0.001):
        if optimizer is None:
            self.optimizer = torch.optim.Adam(self.parameters(), lr=lr)
        else:
            self.optimizer = optimizer

    def extend_model(self, new_component):
        self.add_module(new_component.name, new_component)
        self.update_forward_pass(new_component)

    def update_forward_pass(self, new_component):
        pass


class MyModel(BaseModel):
    def initialize_common_components(self):
        self.layer1 = nn.Linear(10, 15)
        self.layer2 = nn.ReLU()
        self.layer3 = nn.Linear(15, 10)

    def core_computation(self, input):
        x = self.layer1(input)
        x = self.layer2(x)
        x = self.layer3(x)
        return x


def my_model_function():
    return MyModel()


def GetInput():
    return torch.randn(1, 10)


def run_eager(model, x):
    model.eval()
    with torch.no_grad():
        return model(x)


def run_compiled(model, x):
    model.eval()
    compiled_model = torch.compile(
        model,
        fullgraph=True,
        dynamic=True,
        backend="inductor",
    )
    with torch.no_grad():
        return compiled_model(x)


def main():
    torch.manual_seed(0)
    print("torch version:", torch.__version__)

    model = my_model_function()
    x = GetInput()

    print("input shape:", tuple(x.shape))
    print("input dtype:", x.dtype)
    print("model class:", model.__class__.__name__)

    eager_out = run_eager(model, x.clone())
    print("\n[EAGER] run succeeded.")
    print("[EAGER] output shape:", tuple(eager_out.shape))
    print("[EAGER] output dtype:", eager_out.dtype)
    print("[EAGER] output:", eager_out)

    compiled_out = run_compiled(model, x.clone())
    print(compiled_out)


if __name__ == "__main__":
    main()

---

PyTorch version:  2.12.0a0+git775500a
Is debug build: True
CUDA used to build PyTorch: 12.6
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.4 LTS (x86_64)
GCC version: (Ubuntu 12.3.0-1ubuntu1~22.04) 12.3.0
Clang version: Could not collect
CMake version: version 3.22.1
Libc version: glibc-2.35

Python version: 3.10.16 (main, Dec 11 2024, 16:24:50) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-6.8.0-59-generic-x86_64-with-glibc2.35
Is CUDA available: True