Skip to content

Commit

Permalink
Forward OptimizedModule.__setattr__ to the wrapped module (#122098)
Browse files Browse the repository at this point in the history
Summary:
Fixes #114844

In the linked issue we have
```
compiled_module = torch.compile(module)
compiled_module.x = ...
compiled_module(...)  # Mutates self.x
```
Where since the module mutates `self.x` you would expect `compiled_module.x`
to be updated but actually `compiled_module.x = ...` sets an attribute "x"
on the `OptimizedModule` object while the forward method of the module mutates
`module.x`.

This gives the expected behavior by forwarding `compiled_module.__setattr__`
down to `module.__setattr__`. There is already a corresponding `__getattr__`
so now `compiled_module.x` becomes an alias for `module.x`.

X-link: pytorch/pytorch#122098
Approved by: https://github.com/ezyang, https://github.com/lezcano

Reviewed By: clee2000

Differential Revision: D55617139

fbshipit-source-id: afbe861a773128c818e1af03df2b0eeb2a645238

Co-authored-by: Edward Z. Yang <ezyang@meta.com>
  • Loading branch information
2 people authored and facebook-github-bot committed Apr 2, 2024
1 parent 95f31c1 commit da2a8ba
Showing 1 changed file with 0 additions and 10 deletions.
10 changes: 0 additions & 10 deletions userbenchmark/dynamo/dynamobench/_dynamo/testing.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,16 +43,6 @@ def clone_me(x):
return x.detach().clone().requires_grad_(x.requires_grad)


def named_parameters_for_optimized_module(mod):
assert isinstance(mod, eval_frame.OptimizedModule)
return mod._orig_mod.named_parameters


def named_buffers_for_optimized_module(mod):
assert isinstance(mod, eval_frame.OptimizedModule)
return mod._orig_mod.named_buffers


def remove_optimized_module_prefix(name) -> str:
return re.sub(r"^_orig_mod[.]", "", name)

Expand Down

0 comments on commit da2a8ba

Please sign in to comment.