Reland of D65489998 Optimize sharding performance of embeddings" (#3549)

Summary:
X-link: pytorch/torchrec#2664

Pull Request resolved: #3549

X-link: facebookresearch/FBGEMM#634

This diff is a reland of D65489998 after backout in D66800554.

Reviewed By: iamzainhuda

Differential Revision: D66828907

fbshipit-source-id: ab6e6a9faa8255c4847a69a8efb46182bedc9737

fbgemm_gpu/fbgemm_gpu/tbe/ssd/utils/partially_materialized_tensor.py

@@ -48,7 +48,7 @@ def __init__(self, wrapped) -> None:
             wrapped: torch.classes.fbgemm.KVTensorWrapper
         """
         self._wrapped = wrapped
-        self._requires_grad = True
+        self._requires_grad = False
 
     @property
     def wrapped(self):