pytorch

dc18048d - [PT-D][Fix] Broken sharded embedding and embedding bag test fix (#69725)

[PT-D][Fix] Broken sharded embedding and embedding bag test fix (#69725) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/69725 We have added a `no_grad` cx manager in the tensor sharding to ensure that the local_shard is the root node. But it turns out for embedding and embedding_bag, when the `max_norm` is specified, it will complain for row-wise sharding. We use the original `max_norm` of the operators. Error traces: ``` File "/data/sandcastle/boxes/fbsource/fbcode/buck-out/dev/gen/caffe2/test/distributed/_sharded_tensor/sharded_embedding#binary,link-tree/torch/overrides.py", line 1389, in handle_torch_function result = torch_func_method(public_api, types, args, kwargs) File "/data/sandcastle/boxes/fbsource/fbcode/buck-out/dev/gen/caffe2/test/distributed/_sharded_tensor/sharded_embedding#binary,link-tree/torch/distributed/_sharded_tensor/api.py", line 554, in __torch_function__ return sharded_embedding(types, args, kwargs, self._process_group) File "/data/sandcastle/boxes/fbsource/fbcode/buck-out/dev/gen/caffe2/test/distributed/_sharded_tensor/sharded_embedding#binary,link-tree/torch/distributed/_sharded_tensor/ops/embedding.py", line 115, in sharded_embedding return _handle_row_wise_sharding( File "/data/sandcastle/boxes/fbsource/fbcode/buck-out/dev/gen/caffe2/test/distributed/_sharded_tensor/sharded_embedding#binary,link-tree/torch/distributed/_sharded_tensor/ops/embedding.py", line 309, in _handle_row_wise_sharding gathered_input_embeddings = torch.nn.functional.embedding( File "/data/sandcastle/boxes/fbsource/fbcode/buck-out/dev/gen/caffe2/test/distributed/_sharded_tensor/sharded_embedding#binary,link-tree/torch/nn/functional.py", line 2153, in embedding return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse) RuntimeError: A view was created in no_grad mode and its base or another view of its base has been modified inplace with grad mode enabled. Given that this use case is ambiguous and error-prone, it is forbidden. You can clarify your code by moving both the view and the inplace either both inside the no_grad block (if you don't want the inplace to be tracked) or both outside (if you want the inplace to be tracked). exiting process 2 with exit code: 10 ``` As a fix, we clone, detach the local shard from the narrow result without using the context manager. ghstack-source-id: 145773748 Test Plan: CI + Unit test. Reviewed By: pritamdamania87, wanchaol Differential Revision: D33000927 fbshipit-source-id: 4d5a93120675e90d4d6d6225a51c4a481d18d159