Fix empty EmbeddingCollection/EmbeddingBagCollection edge cases (#2823)

Summary:
Pull Request resolved: #2823

This is an edge-case fix for D68991644. Some models include empty EmbeddingCollections or EmbeddingBagCollections by mistake (?), such as T217919588, thus causing the feature order caching logic broken.

Reviewed By: sidt-meta

Differential Revision: D71218553

fbshipit-source-id: 204661b4327cd792e3e98f482126d9634f93b011

Author: joshuuuasu

torchrec/quant/embedding_modules.py

@@ -358,7 +358,7 @@ def __init__(
         self._feature_names: List[str] = []
         self._feature_splits: List[int] = []
         self._length_per_key: List[int] = []
-        self._features_order: List[int] = []
+        self._features_order: Optional[List[int]] = None
         # Registering in a List instead of ModuleList because we want don't want them to be auto-registered.
         # Their states will be modified via self.embedding_bags
         self._emb_modules: List[nn.Module] = []
@@ -502,20 +502,22 @@ def forward(
         kjt_keys = _get_kjt_keys(features)
         # Cache the features order since the features will always have the same order of keys in inference.
         if getattr(self, MODULE_ATTR_CACHE_FEATURES_ORDER, False):
-            if self._features_order == []:
-                for k in self._feature_names:
-                    self._features_order.append(kjt_keys.index(k))
-                self.register_buffer(
-                    "_features_order_tensor",
-                    torch.tensor(
-                        data=self._features_order,
-                        device=features.device(),
-                        dtype=torch.int32,
-                    ),
-                    persistent=False,
-                )
+            if self._features_order is None:
+                self._features_order = [kjt_keys.index(k) for k in self._feature_names]
+                if self._features_order:
+                    self.register_buffer(
+                        "_features_order_tensor",
+                        torch.tensor(
+                            data=self._features_order,
+                            device=features.device(),
+                            dtype=torch.int32,
+                        ),
+                        persistent=False,
+                    )
             kjt_permute = _permute_kjt(
-                features, self._features_order, self._features_order_tensor
+                features,
+                self._features_order,
+                getattr(self, "_features_order_tensor", None),
             )
         else:
             kjt_permute_order = [kjt_keys.index(k) for k in self._feature_names]
@@ -752,7 +754,7 @@ def __init__(  # noqa C901
         self.row_alignment = row_alignment
         self._key_to_tables: Dict[DataType, List[EmbeddingConfig]] = defaultdict(list)
         self._feature_names: List[str] = []
-        self._features_order: List[int] = []
+        self._features_order: Optional[List[int]] = None
 
         self._table_name_to_quantized_weights: Optional[
             Dict[str, Tuple[Tensor, Tensor]]
@@ -881,20 +883,22 @@ def forward(
         kjt_keys = _get_kjt_keys(features)
         # Cache the features order since the features will always have the same order of keys in inference.
         if getattr(self, MODULE_ATTR_CACHE_FEATURES_ORDER, False):
-            if self._features_order == []:
-                for k in self._feature_names:
-                    self._features_order.append(kjt_keys.index(k))
-                self.register_buffer(
-                    "_features_order_tensor",
-                    torch.tensor(
-                        data=self._features_order,
-                        device=features.device(),
-                        dtype=torch.int32,
-                    ),
-                    persistent=False,
-                )
+            if self._features_order is None:
+                self._features_order = [kjt_keys.index(k) for k in self._feature_names]
+                if self._features_order:
+                    self.register_buffer(
+                        "_features_order_tensor",
+                        torch.tensor(
+                            data=self._features_order,
+                            device=features.device(),
+                            dtype=torch.int32,
+                        ),
+                        persistent=False,
+                    )
             kjt_permute = _permute_kjt(
-                features, self._features_order, self._features_order_tensor
+                features,
+                self._features_order,
+                getattr(self, "_features_order_tensor", None),
             )
         else:
             kjt_permute_order = [kjt_keys.index(k) for k in self._feature_names]