Enable MiscTests.test_map_side_effects to work with nn module inlining

[laithsakka]

TORCHDYNAMO_INLINE_INBUILT_NN_MODULES=1 python test/dynamo/test_misc.py -k MiscTests.test_map_side_effects
error:

You can suppress this exception and fall back to eager by setting:
    import torch._dynamo
    torch._dynamo.config.suppress_errors = True


During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/data/users/lsakka/pytorch/pytorch/torch/testing/_internal/common_utils.py", line 2756, in wrapper
    method(*args, **kwargs)
  File "/data/users/lsakka/pytorch/pytorch/test/dynamo/test_misc.py", line 5384, in test_map_side_effects
    with self.assertRaisesRegex(
AssertionError: "Can't inplace modify module params/buffers" does not match "HigherOrderOperator: Mutating a variable not in the current scope (SideEffects)

from user code:
   File "/data/users/lsakka/pytorch/pytorch/test/dynamo/test_misc.py", line 5381, in forward
    return map(body, xs)
  File "/data/users/lsakka/pytorch/pytorch/test/dynamo/test_misc.py", line 5378, in body
    self.w += 1

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information

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

[laithsakka]

related readings:
https://pytorch-dev-podcast.simplecast.com/episodes/higher-order-operators
https://dev-discuss.pytorch.org/t/higher-order-operators-2023-10/1565

[laithsakka]

First anyway its not allowed to mutate the nn module variable from with in the map operation, the error message is different due to the following:
when nn module inlining not enabled, the code in builtin.py

     elif isinstance(obj, variables.NNModuleVariable):
            if not tx.output.is_root_tracer():
                raise AttributeMutationError(
                    "Can't inplace modify module params/buffers inside HigherOrderOp"
                )

is triggered, this is not triggered when inlining is enabled because obj in that case is UnspecializedNNModuleVariable.
instead what happen is that this check is by passed and later on another check is triggered when set_attribute is called
on UserDefinedObjectVariable:
tx.output.side_effects.store_attr(self, name, value)
-->

       # People do things like self.dim = dim inside autograd.Function.
        # These are benign.
        if isinstance(item, AutogradFunctionContextVariable):
            return True
        if not is_side_effect_safe(item.mutable_local):
            unimplemented(
                "HigherOrderOperator: Mutating a variable not in the current scope (SideEffects)"
            )

[laithsakka]

We have two options here:

1. do not allow mutations even for UnspecializedNNModuleVariable.
2. change the unit test to expect different messages if inlining is enabled vs disabled,
   or just check that unsupported is thrown.
   @anijain2305 sounds like (2) is the appropriate let me know otherwise.

[ezyang]

IMO it should be possible for Dynamo to trace through NN module mutations

[laithsakka]

IMO it should be possible for Dynamo to trace through NN module mutations
it does in general, but it does not here when it that happens from within a ma high order operation.
ex: (this work)

class Module(torch.nn.Module):
   def __init__(self):
           super().__init__()
           self.w = torch.tensor(1)
           
    def forward(self, xs):
            def body(x):
                    self.w += 1
                    return x

                return body(xs)

ex2: this does not work

class Module(torch.nn.Module):
  def __init__(self):
          super().__init__()
          self.w = torch.tensor(1)
 def forward(self, xs):
               def body(x):
                   self.w += 1
                   return x

               return map(body, xs)

@ezyang do you think we should make ex2 also work above ? it is treated with inlining enabled as a UserDefinedObjectVariable

[ezyang]

If map is the hop, no I don't expect it to work, but this should have failed without nn module inlining too :P

[laithsakka]

If map is the hop, no I don't expect it to work, but this should have failed without nn module inlining too :P

yep it is failing both ways. Just the error message changed