[guards][cpp-guards] Optimize NN module getattr guards

[anijain2305]

Stack from ghstack (oldest at bottom):

• [dynamo] Add ID_MATCH guards on inlined functions to force compilation on monkeypatching #124975
• [dynamo] Turn on guard_nn_modules #125202
• [dynamo][prepare for nn module guards] Guard nn modules for a few benchmarks #125324
• -> [guards][cpp-guards] Optimize NN module getattr guards #124522
• [dynamo][nn module] Check for duplicate tensors in register_attr_or_module #125421
• [dynamo][cpp-guards] Improve recompilation reason logic for NO_TENSOR_ALIASING guard #125439

Improves the guard overhead of MobileBert model with nn module guards from 92000 units to 20000 units.

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @fegin @XilunWu @wanchaol @fduwjj @wz337 @tianyu-l @wconstab @yf225 @chauhang @d4l3k @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/124522

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (1 Unrelated Failure)

As of commit 4af936f with merge base 3b5f6b1 ():

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• slow / linux-jammy-py3.10-clang15-asan / test (slow, 2, 2, linux.4xlarge) (gh) (similar failure)
  !!!!!!!!!!!!!!!!!!!!!!!!!!!!!! KeyboardInterrupt !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[jansel]

add some more tests for this

[jansel]

Are there tests covering this case? (Where we change the order)

[anijain2305]

There were none for NNModuleVariableTracker. There are for UnSpecializedNNModuleVariable.

Good catch. Added and made the fixes as well.

[jansel]

class attributes are shadowed by module dict attributes... I think the order of these are wrong. add some tests for shadowing.

[anijain2305]

You are right .. Thanks for pointing this out, I will fix it. But there is a corner case with submodules, where class attr take precedence. This is the example

import torch
import torch._dynamo.testing

torch._dynamo.config.guard_nn_modules = True


class Mod(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.a = 3

    def forward(self,x, c=4):
        return x * c

    def linear(self, x):
        return x

    # b is class attribute and this is a bug that jansel pointed out
    # mod.b overshadows Mod.b
    def b(self, x):
        assert False
        return x

class MyMod(Mod):
    def __init__(self):
        super().__init__()
        # IF THIS LINE ACTUALLY CHAGES LINEAR, IT WOULD FAIL WITH SHAPE MISMATCH ERROR
        self.linear = torch.nn.Linear(11, 11)
        self.a = 2
        self.b = 2


    def forward(self, x, c=None):
        return self.linear(x) * self.a * self.b


mod = MyMod()
mod(torch.ones(3, 3))

# Instance attributes overwrite the instance attributes of the parent class.
assert mod.a == 2

# ANOMALY - mod.linear is not the parameter
assert mod.linear.__code__ is Mod.linear.__code__

# jansel's valid comment - module dict attr overshadow the class attrs
# anijain2305 - This does not cause a bug because today we fallback to getattr,
# but we are losing perf opportunity here.
assert mod.b == 2

cnts = torch._dynamo.testing.CompileCounter()
opt_mod = torch.compile(mod, backend=cnts)
opt_mod(torch.ones(3, 3))
opt_mod(torch.ones(3, 3))

assert cnts.frame_count == 1

The discrepancy is for mod.linear. Parent class has a function named linear. When the derived class instance calls self.linear, it calls the parent class function.

This is a special case. Apart from that everything else conforms to the general rule - "class attributes are shadowed by the mod dict attributes".

I also think this is a bug in __setattr__ of nn.Module. For example, I am unable to assign a self.param if the parent class already has a function named param. If you think, this is just a bug, I can send a PR for nn modules.

[anijain2305]

@jansel just a gentle ping in case this slipped in your notifications.