[DDP] Fix when buffers are reassigned in module (#64472)
authorRohan Varma <rvarm1@fb.com>
Thu, 9 Sep 2021 02:13:33 +0000 (19:13 -0700)
committerFacebook GitHub Bot <facebook-github-bot@users.noreply.github.com>
Thu, 9 Sep 2021 02:14:55 +0000 (19:14 -0700)
commitd59ecc02df70bad2273858c2fad2b4993133a3d3
treebec62ab7769445c12fa39d9c955c70aec0b5eaec
parentb6544ef81511cc527e4945a969d8a680d2d5069f
[DDP] Fix when buffers are reassigned in module (#64472)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/64472

Sometimes, user module can reassign tensor buffer, as in:

```
self.buffer = torch.randn(1, 2) # in init
self.buffer += 1 # in forward
```

in this case, `self.modules_buffers` will become outdated and we should
repopulate self.modules_buffers if we need to sync module buffers.

See https://github.com/pytorch/pytorch/issues/63916 for full description of the
issue.
ghstack-source-id: 137526309

Test Plan: CI

Reviewed By: zhaojuanmao

Differential Revision: D30745921

fbshipit-source-id: 25eb1edbf445703a481802e07f3058d38ea6fc64
torch/nn/parallel/distributed.py
torch/testing/_internal/distributed/distributed_test.py