dont persist feature order buffers as they dont need to be checkpoint…

…ed (pytorch#956)

Summary:
Pull Request resolved: pytorch#956

TorchRec don't need to persist feature order buffers -> these are correct on instantiation and don't need to be check pointed

Previously, this issue was hidden bc shardedModules buffers() (via embedding kernels buffer call) yielded nothing, which is no longer true as of D41964643 (pytorch@e8ab2de)

Reviewed By: YLGH

Differential Revision: D42591693

fbshipit-source-id: 6694239c54fbbaeb563b1c0221f4a5324f2c99b6

torchrec/distributed/embedding.py

@@ -585,6 +585,7 @@ def _create_input_dist(
         self.register_buffer(
             "_features_order_tensor",
             torch.tensor(self._features_order, device=self._device, dtype=torch.int32),
+            persistent=False,
         )
 
     def _create_lookups(self) -> None:

torchrec/distributed/embedding_tower_sharding.py

@@ -197,6 +197,7 @@ def _create_input_dist(
                 torch.tensor(
                     self._kjt_features_order, device=self._device, dtype=torch.int32
                 ),
+                persistent=False,
             )
 
         if self._wkjt_feature_names != wkjt_feature_names:
@@ -208,6 +209,7 @@ def _create_input_dist(
                 torch.tensor(
                     self._wkjt_features_order, device=self._device, dtype=torch.int32
                 ),
+                persistent=False,
             )
 
         node_count = dist.get_world_size(self._cross_pg)

torchrec/distributed/embeddingbag.py

@@ -542,6 +542,7 @@ def _create_input_dist(
                 torch.tensor(
                     self._features_order, device=self._device, dtype=torch.int32
                 ),
+                persistent=False,
             )
 
     def _create_lookups(

torchrec/distributed/quant_embedding.py

@@ -237,6 +237,7 @@ def _create_input_dist(
         self.register_buffer(
             "_features_order_tensor",
             torch.tensor(self._features_order, device=device, dtype=torch.int32),
+            persistent=False,
         )
 
     def _create_lookups(self, fused_params: Optional[Dict[str, Any]]) -> None: