use qualname for the serialized nn.Module (meta-pytorch#3021)

Summary:
Pull Request resolved: meta-pytorch#3021

# context
* original the IR serializer use type(m).__name__ as the key for the module_to_serializer_cls mapping
* this becomes a problem when there are multiple modules with same name
* this change uses the full qualname of a class as the key
* example
```
# previously
ebc = EmbeddingBagCollection(...)
type(ebc).__name__ == 'EmbeddingBagCollection'

# now
ebc = EmbeddingBagCollection(...)
qualname(ebc) == 'torchrec.modules.embedding_modules.EmbeddingBagCollection'
```
# debug print
> {'torchrec.modules.embedding_modules.EmbeddingBagCollection': <class 'torchrec.ir.serializer.EBCJsonSerializer'>, 'torchrec.modules.feature_processor_.PositionWeightedModule': <class 'torchrec.ir.serializer.PWMJsonSerializer'>, 'torchrec.modules.feature_processor_.PositionWeightedModuleCollection': <class 'torchrec.ir.serializer.PWMCJsonSerializer'>, 'torchrec.modules.fp_embedding_modules.FeatureProcessedEmbeddingBagCollection': <class 'torchrec.ir.serializer.FPEBCJsonSerializer'>, 'torchrec.modules.regroup.KTRegroupAsDict': <class 'torchrec.ir.serializer.KTRegroupAsDictJsonSerializer'>}

# backward compatibility issue
* a previously serialized model (module) won't be correctly deserialized
* the workaround is to keep the old key name in the `module_to_serializer_cls` map for a while until the old method is fully deprecated

Differential Revision: D75727469

Author: TroyGarden

torchrec/ir/serializer.py

@@ -22,7 +22,7 @@
 )
 
 from torchrec.ir.types import SerializerInterface
-from torchrec.ir.utils import logging
+from torchrec.ir.utils import logging, qualname
 from torchrec.modules.embedding_configs import DataType, EmbeddingBagConfig, PoolingType
 from torchrec.modules.embedding_modules import EmbeddingBagCollection
 from torchrec.modules.feature_processor_ import (
@@ -174,7 +174,7 @@ def swap_meta_forward(cls, module: nn.Module) -> None:
 
     @classmethod
     def encapsulate_module(cls, module: nn.Module) -> List[str]:
-        typename = type(module).__name__
+        typename = qualname(module)
         serializer = cls.module_to_serializer_cls.get(typename)
         if serializer is None:
             raise ValueError(
@@ -184,8 +184,8 @@ def encapsulate_module(cls, module: nn.Module) -> List[str]:
         assert serializer._module_cls is not None
         if not isinstance(module, serializer._module_cls):
             raise ValueError(
-                f"Expected module to be of type {serializer._module_cls.__name__}, "
-                f"got {type(module)}"
+                f"Expected module to be of type {qualname(serializer._module_cls)}, "
+                f"got {qualname(module)}"
             )
         metadata_dict = serializer.serialize_to_dict(module)
         raw_dict = {"typename": typename, "metadata_dict": metadata_dict}
@@ -218,7 +218,7 @@ def decapsulate_module(
             )
         if not isinstance(module, serializer._module_cls):
             raise ValueError(
-                f"Expected module to be of type {serializer._module_cls.__name__}, got {type(module)}"
+                f"Expected module to be of type {qualname(serializer._module_cls)}, got {qualname(module)}"
             )
         return module
 
@@ -275,7 +275,9 @@ def deserialize_from_dict(
         )
 
 
-JsonSerializer.module_to_serializer_cls["EmbeddingBagCollection"] = EBCJsonSerializer
+JsonSerializer.module_to_serializer_cls[qualname(EmbeddingBagCollection)] = (
+    EBCJsonSerializer
+)
 
 
 class PWMJsonSerializer(JsonSerializer):
@@ -299,7 +301,9 @@ def deserialize_from_dict(
         return PositionWeightedModule(metadata.max_feature_length, device)
 
 
-JsonSerializer.module_to_serializer_cls["PositionWeightedModule"] = PWMJsonSerializer
+JsonSerializer.module_to_serializer_cls[qualname(PositionWeightedModule)] = (
+    PWMJsonSerializer
+)
 
 
 class PWMCJsonSerializer(JsonSerializer):
@@ -331,7 +335,7 @@ def deserialize_from_dict(
         return PositionWeightedModuleCollection(max_feature_lengths, device)
 
 
-JsonSerializer.module_to_serializer_cls["PositionWeightedModuleCollection"] = (
+JsonSerializer.module_to_serializer_cls[qualname(PositionWeightedModuleCollection)] = (
     PWMCJsonSerializer
 )
 
@@ -385,9 +389,9 @@ def deserialize_from_dict(
         )
 
 
-JsonSerializer.module_to_serializer_cls["FeatureProcessedEmbeddingBagCollection"] = (
-    FPEBCJsonSerializer
-)
+JsonSerializer.module_to_serializer_cls[
+    qualname(FeatureProcessedEmbeddingBagCollection)
+] = FPEBCJsonSerializer
 
 
 class KTRegroupAsDictJsonSerializer(JsonSerializer):
@@ -436,6 +440,6 @@ def deserialize_from_dict(
         )
 
 
-JsonSerializer.module_to_serializer_cls["KTRegroupAsDict"] = (
+JsonSerializer.module_to_serializer_cls[qualname(KTRegroupAsDict)] = (
     KTRegroupAsDictJsonSerializer
 )

torchrec/ir/tests/test_serializer.py

@@ -21,6 +21,7 @@
     decapsulate_ir_modules,
     encapsulate_ir_modules,
     mark_dynamic_kjt,
+    qualname,
 )
 from torchrec.modules.embedding_configs import EmbeddingBagConfig
 from torchrec.modules.embedding_modules import EmbeddingBagCollection
@@ -524,7 +525,7 @@ def forward(self, features: KeyedJaggedTensor) -> List[torch.Tensor]:
 
         eager_out = model(id_list_features)
 
-        JsonSerializer.module_to_serializer_cls["CompoundModule"] = (
+        JsonSerializer.module_to_serializer_cls[qualname(CompoundModule)] = (
             CompoundModuleSerializer
         )
         # Serialize

torchrec/ir/utils.py

@@ -12,7 +12,7 @@
 import logging
 import operator
 from collections import defaultdict
-from typing import Dict, List, Optional, Tuple, Type
+from typing import Dict, List, Optional, Tuple, Type, Union
 
 import torch
 
@@ -34,6 +34,13 @@
 logger: logging.Logger = logging.getLogger(__name__)
 
 
+def qualname(m: Union[nn.Module, type[nn.Module]]) -> str:
+    if isinstance(m, nn.Module):
+        return type(m).__module__ + "." + type(m).__qualname__
+    else:
+        return m.__module__ + "." + m.__qualname__
+
+
 def get_device(tensors: List[Optional[torch.Tensor]]) -> Optional[torch.device]:
     """
     Returns the device of the first non-None tensor in the list.
@@ -115,7 +122,7 @@ def encapsulate_ir_modules(
     preserve_fqns: List[str] = []  # fqns of the serialized modules
     children: List[str] = []  # fqns of the children that need further serialization
     # handle current module, and find the children which need further serialization
-    if type(module).__name__ in serializer.module_to_serializer_cls:
+    if qualname(module) in serializer.module_to_serializer_cls:
         children = serializer.encapsulate_module(module)
         preserve_fqns.append(fqn)
     else: