refactor using encapsulate and decapsulate (#2174)

Summary:
Pull Request resolved: #2174

# context
1. we plan to consolidate the `non_strict_forward` into the TorchRec IR serializer, and this diff is to implement the new set of APIs for testing and migration.
**a.** current TorchRec IR APIs are `serialize_embedding_modules` and `deserialize_embedding_modules`.
**b.** the new set of APIs are named `encapsulate_ir_modules` and `decapsulate_ir_modules`.
**c.** two sets of APIs basically take in the same arguments.
3. the major differences are that:
**a.** in the new `encapsulate_ir_modules`, the TorchRec (embedding) modules' forward functions are replaced by a `ir_meta_forward` via the corresponding serializer. While in the old API flow the `non_strict_forward` is embedded within the TorchRec module and hidden behind a `is_non_strict_export` flag.
**b.** in the new `decapsulate_ir_modules`, it takes in the `unflatten_module` which is unflattened by `torch.export.unflatten`. While in the old API, it takes in the `ExportedProgram` and run the `torch.export.unflatten` function inside the TorchRec IR's `deserialize_embedding_modules` function, which introduced extra coupling.
3. This diff does **NOT** affect the original APIs.
4. Embedding modules including EBC, PEA, VLE, fpEBC have been tested with dynamic shape support.

# details
* schema
 {F1733431041}  {F1733431375}

* static ir_custom_op definition, **dynamic shape is supported**
```
torch.library.custom_op("torchrec::ir_custom_op", mutates_args={})
def ir_custom_op_impl(
    tensors: List[Optional[torch.Tensor]], batch_size: int, dim: int
) -> torch.Tensor: # when multiple output tensor is needed, we can simply do a torch.split
    device = None
    for t in tensors:
        if t is not None:
            device = t.device
            break
    logger.info(f"torch.ops.torchrec.ir_custom_op -> ({batch_size}, {dim})")
    return torch.empty(batch_size, dim, device=device)
```
* The previous `custom_op` is unified into a static operator named `ir_custom_op`, which does not require registering on the fly. So there is no need to call a helper function `register_custom_ops_for_nodes` when unflatten an exported IR artifact.
* It decouples the `non_strict_export_forward` function in each TorchRec module, and avoids the complexity in the original module's forward logic as below, and the `non_strict_export` flag in `torch.export.export` is no longer a necessary flag.
```
    def forward(
        self,
        id_list_features: KeyedJaggedTensor,
        id_score_list_features: Optional[KeyedJaggedTensor] = None,
    ) -> Dict[str, torch.Tensor]:
        if is_non_strict_exporting() and not torch.jit.is_scripting():
            return self._non_strict_exporting_forward(  # <---- this create a shortcut to the original forward function
                id_list_features, id_score_list_features
            )
        ... # the actual forward function
```

Reviewed By: dstaay-fb, PaulZhang12

Differential Revision: D59019375

fbshipit-source-id: ef0e539a7e0d11a206fb62485071e9e9cf2887bc

Author: TroyGarden

torchrec/ir/serializer.py

@@ -8,7 +8,6 @@
 # pyre-strict
 
 import json
-import logging
 from typing import Any, Dict, List, Optional, Tuple, Type
 
 import torch
@@ -23,6 +22,7 @@
 )
 
 from torchrec.ir.types import SerializerInterface
+from torchrec.ir.utils import logging
 from torchrec.modules.embedding_configs import DataType, EmbeddingBagConfig, PoolingType
 from torchrec.modules.embedding_modules import EmbeddingBagCollection
 from torchrec.modules.feature_processor_ import (
@@ -32,6 +32,8 @@
     PositionWeightedModuleCollection,
 )
 from torchrec.modules.fp_embedding_modules import FeatureProcessedEmbeddingBagCollection
+from torchrec.sparse.jagged_tensor import KeyedJaggedTensor, KeyedTensor
+
 
 logger: logging.Logger = logging.getLogger(__name__)
 
@@ -82,6 +84,26 @@ def get_deserialized_device(
     return device
 
 
+def ebc_meta_forward(
+    ebc: EmbeddingBagCollection,
+    features: KeyedJaggedTensor,
+) -> KeyedTensor:
+    batch_size = features.stride()
+    dim = sum(ebc._lengths_per_embedding)
+    arg_list = [
+        features.values(),
+        features.weights_or_none(),
+        features.lengths_or_none(),
+        features.offsets_or_none(),
+    ]  # if want to include the weights: `+ [bag.weight for bag in self.embedding_bags.values()]`
+    output = torch.ops.torchrec.ir_custom_op(arg_list, batch_size, dim)
+    return KeyedTensor(
+        keys=ebc._embedding_names,
+        values=output,
+        length_per_key=ebc._lengths_per_embedding,
+    )
+
+
 class JsonSerializer(SerializerInterface):
     """
     Serializer for torch.export IR using json.
@@ -163,10 +185,70 @@ def deserialize(
             )
         return module
 
+    @classmethod
+    def swap_meta_forward(cls, module: nn.Module) -> None:
+        pass
+
+    @classmethod
+    def encapsulate_module(cls, module: nn.Module) -> List[str]:
+        typename = type(module).__name__
+        serializer = cls.module_to_serializer_cls.get(typename)
+        if serializer is None:
+            raise ValueError(
+                f"Expected typename to be one of {list(cls.module_to_serializer_cls.keys())}, got {typename}"
+            )
+        assert issubclass(serializer, JsonSerializer)
+        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)}"
+            )
+        metadata_dict = serializer.serialize_to_dict(module)
+        raw_dict = {"typename": typename, "metadata_dict": metadata_dict}
+        ir_metadata_tensor = torch.frombuffer(
+            json.dumps(raw_dict).encode(), dtype=torch.uint8
+        )
+        module.register_buffer("ir_metadata", ir_metadata_tensor, persistent=False)
+        serializer.swap_meta_forward(module)
+        return serializer.children(module)
+
+    @classmethod
+    def decapsulate_module(
+        cls, module: nn.Module, device: Optional[torch.device] = None
+    ) -> nn.Module:
+        raw_bytes = module.get_buffer("ir_metadata").numpy().tobytes()
+        raw_dict = json.loads(raw_bytes.decode())
+        typename = raw_dict["typename"]
+        metadata_dict = raw_dict["metadata_dict"]
+        if typename not in cls.module_to_serializer_cls:
+            raise ValueError(
+                f"Expected typename to be one of {list(cls.module_to_serializer_cls.keys())}, got {typename}"
+            )
+        serializer = cls.module_to_serializer_cls[typename]
+        assert issubclass(serializer, JsonSerializer)
+        module = serializer.deserialize_from_dict(metadata_dict, device, module)
+
+        if serializer._module_cls is None:
+            raise ValueError(
+                "Must assign a nn.Module to class static variable _module_cls"
+            )
+        if not isinstance(module, serializer._module_cls):
+            raise ValueError(
+                f"Expected module to be of type {serializer._module_cls.__name__}, got {type(module)}"
+            )
+        return module
+
 
 class EBCJsonSerializer(JsonSerializer):
     _module_cls = EmbeddingBagCollection
 
+    @classmethod
+    def swap_meta_forward(cls, module: nn.Module) -> None:
+        assert isinstance(module, cls._module_cls)
+        # pyre-ignore
+        module.forward = ebc_meta_forward.__get__(module, cls._module_cls)
+
     @classmethod
     def serialize_to_dict(
         cls,

torchrec/ir/tests/test_serializer.py

@@ -18,9 +18,9 @@
 from torchrec.ir.serializer import JsonSerializer
 
 from torchrec.ir.utils import (
-    deserialize_embedding_modules,
+    decapsulate_ir_modules,
+    encapsulate_ir_modules,
     mark_dynamic_kjt,
-    serialize_embedding_modules,
 )
 
 from torchrec.modules.embedding_configs import EmbeddingBagConfig
@@ -30,7 +30,6 @@
     PositionWeightedModuleCollection,
 )
 from torchrec.modules.fp_embedding_modules import FeatureProcessedEmbeddingBagCollection
-from torchrec.modules.utils import operator_registry_state
 from torchrec.sparse.jagged_tensor import KeyedJaggedTensor, KeyedTensor
 
 
@@ -183,7 +182,7 @@ def test_serialize_deserialize_ebc(self) -> None:
         eager_out = model(id_list_features)
 
         # Serialize EBC
-        model, sparse_fqns = serialize_embedding_modules(model, JsonSerializer)
+        model, sparse_fqns = encapsulate_ir_modules(model, JsonSerializer)
         ep = torch.export.export(
             model,
             (id_list_features,),
@@ -199,34 +198,9 @@ def test_serialize_deserialize_ebc(self) -> None:
         for i, tensor in enumerate(ep_output):
             self.assertEqual(eager_out[i].shape, tensor.shape)
 
-        # Should have 3 custom op registered, as dimensions of ebc are same,
-        # and two fpEBCs have different dims
-        self.assertEqual(len(operator_registry_state.op_registry_schema), 3)
-
-        total_dim_ebc = sum(model.ebc1._lengths_per_embedding)
-        total_dim_fpebc1 = sum(
-            model.fpebc1._embedding_bag_collection._lengths_per_embedding
-        )
-        total_dim_fpebc2 = sum(
-            model.fpebc2._embedding_bag_collection._lengths_per_embedding
-        )
-        # Check if custom op is registered with the correct name
-        # EmbeddingBagCollection type and total dim
-        self.assertTrue(
-            f"EmbeddingBagCollection_{total_dim_ebc}"
-            in operator_registry_state.op_registry_schema
-        )
-        self.assertTrue(
-            f"EmbeddingBagCollection_{total_dim_fpebc1}"
-            in operator_registry_state.op_registry_schema
-        )
-        self.assertTrue(
-            f"EmbeddingBagCollection_{total_dim_fpebc2}"
-            in operator_registry_state.op_registry_schema
-        )
-
         # Deserialize EBC
-        deserialized_model = deserialize_embedding_modules(ep, JsonSerializer)
+        unflatten_ep = torch.export.unflatten(ep)
+        deserialized_model = decapsulate_ir_modules(unflatten_ep, JsonSerializer)
 
         # check EBC config
         for i in range(5):
@@ -265,9 +239,8 @@ def test_serialize_deserialize_ebc(self) -> None:
                 self.assertEqual(deserialized.num_embeddings, orginal.num_embeddings)
                 self.assertEqual(deserialized.feature_names, orginal.feature_names)
 
-        deserialized_model.load_state_dict(model.state_dict())
-
         # Run forward on deserialized model and compare the output
+        deserialized_model.load_state_dict(model.state_dict())
         deserialized_out = deserialized_model(id_list_features)
 
         self.assertEqual(len(deserialized_out), len(eager_out))
@@ -292,7 +265,7 @@ def test_dynamic_shape_ebc(self) -> None:
 
         # Serialize EBC
         collection = mark_dynamic_kjt(feature1)
-        model, sparse_fqns = serialize_embedding_modules(model, JsonSerializer)
+        model, sparse_fqns = encapsulate_ir_modules(model, JsonSerializer)
         ep = torch.export.export(
             model,
             (feature1,),
@@ -311,7 +284,8 @@ def test_dynamic_shape_ebc(self) -> None:
             self.assertEqual(eager_out[i].shape, tensor.shape)
 
         # Deserialize EBC
-        deserialized_model = deserialize_embedding_modules(ep, JsonSerializer)
+        unflatten_ep = torch.export.unflatten(ep)
+        deserialized_model = decapsulate_ir_modules(unflatten_ep, JsonSerializer)
         deserialized_model.load_state_dict(model.state_dict())
 
         # Run forward on deserialized model
@@ -330,7 +304,7 @@ def test_deserialized_device(self) -> None:
         )
 
         # Serialize EBC
-        model, sparse_fqns = serialize_embedding_modules(model, JsonSerializer)
+        model, sparse_fqns = encapsulate_ir_modules(model, JsonSerializer)
         ep = torch.export.export(
             model,
             (id_list_features,),
@@ -345,8 +319,9 @@ def test_deserialized_device(self) -> None:
             if device == "cuda" and not torch.cuda.is_available():
                 continue
             device = torch.device(device)
-            deserialized_model = deserialize_embedding_modules(
-                ep, JsonSerializer, device
+            unflatten_ep = torch.export.unflatten(ep)
+            deserialized_model = decapsulate_ir_modules(
+                unflatten_ep, JsonSerializer, device
             )
             for name, m in deserialized_model.named_modules():
                 if hasattr(m, "device"):
@@ -408,7 +383,7 @@ def forward(self, features: KeyedJaggedTensor) -> List[torch.Tensor]:
             CompoundModuleSerializer
         )
         # Serialize
-        model, sparse_fqns = serialize_embedding_modules(model, JsonSerializer)
+        model, sparse_fqns = encapsulate_ir_modules(model, JsonSerializer)
         ep = torch.export.export(
             model,
             (id_list_features,),
@@ -424,8 +399,8 @@ def forward(self, features: KeyedJaggedTensor) -> List[torch.Tensor]:
             self.assertEqual(x.shape, y.shape)
 
         # Deserialize
-        deserialized_model = deserialize_embedding_modules(ep, JsonSerializer)
-
+        unflatten_ep = torch.export.unflatten(ep)
+        deserialized_model = decapsulate_ir_modules(unflatten_ep, JsonSerializer)
         # Check if Compound Module is deserialized correctly
         self.assertIsInstance(deserialized_model.comp, CompoundModule)
         self.assertIsInstance(deserialized_model.comp.comp, CompoundModule)

torchrec/ir/types.py

@@ -47,3 +47,18 @@ def deserialize(
     ) -> nn.Module:
         # Take the bytes in the buffer and regenerate the eager embedding module
         raise NotImplementedError
+
+    @classmethod
+    @abc.abstractmethod
+    def encapsulate_module(cls, module: nn.Module) -> List[str]:
+        # Take the eager embedding module and encapsulate the module, including serialization
+        # and meta_forward-swapping, then returns a list of children (fqns) which needs further encapsulation
+        raise NotImplementedError
+
+    @classmethod
+    @abc.abstractmethod
+    def decapsulate_module(
+        cls, module: nn.Module, device: Optional[torch.device] = None
+    ) -> nn.Module:
+        # Take the eager embedding module and decapsulate it by removing serialization and meta_forward-swapping
+        raise NotImplementedError

torchrec/ir/utils.py

@@ -9,6 +9,7 @@
 
 #!/usr/bin/env python3
 
+import logging
 from collections import defaultdict
 from typing import Dict, List, Optional, Tuple, Type, Union
 
@@ -24,6 +25,83 @@
 # TODO: Replace the default interface with the python dataclass interface
 DEFAULT_SERIALIZER_CLS = SerializerInterface
 DYNAMIC_DIMS: Dict[str, int] = defaultdict(int)
+logger: logging.Logger = logging.getLogger(__name__)
+
+
+@torch.library.custom_op("torchrec::ir_custom_op", mutates_args={})
+def ir_custom_op_impl(
+    tensors: List[Optional[torch.Tensor]], batch_size: int, dim: int
+) -> torch.Tensor:
+    device = None
+    for t in tensors:
+        if t is not None:
+            device = t.device
+            break
+    logger.info(f"torch.ops.torchrec.ir_custom_op -> ({batch_size}, {dim})")
+    return torch.empty(batch_size, dim, device=device)
+
+
+@torch.library.register_fake("torchrec::ir_custom_op")
+def ir_custom_op_fake(
+    tensors: List[Optional[torch.Tensor]], batch_size: int, dim: int
+) -> torch.Tensor:
+    logger.info(f"ir_custom_op_fake -> ({batch_size}, {dim})")
+    return torch.empty(batch_size, dim)
+
+
+def encapsulate_ir_modules(
+    module: nn.Module,
+    serializer: Type[SerializerInterface] = DEFAULT_SERIALIZER_CLS,
+    fqn: str = "",
+) -> Tuple[nn.Module, List[str]]:
+    """
+    Takes a module and encapsulate its embedding modules and serializes them to the module buffer.
+    Returns the modified module and a list of fqns that had the buffer added, which is needed for torch.export
+    The encapsulation is done by using meta_forward function provided by the serializer
+    to replace the module's original forward function.
+    """
+    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:
+        children = serializer.encapsulate_module(module)
+        preserve_fqns.append(fqn)
+    else:
+        # if the module is not of type serializer, then we check all its children
+        children = [child for child, _ in module.named_children()]
+
+    # handle child modules recursively
+    for child in children:
+        submodule = module.get_submodule(child)
+        child_fqn = f"{fqn}.{child}" if len(fqn) > 0 else child
+        _, fqns = encapsulate_ir_modules(submodule, serializer, child_fqn)
+        preserve_fqns.extend(fqns)
+    return module, preserve_fqns
+
+
+def decapsulate_ir_modules(
+    module: nn.Module,
+    serializer: Type[SerializerInterface] = DEFAULT_SERIALIZER_CLS,
+    device: Optional[torch.device] = None,
+) -> nn.Module:
+    """
+    Takes a module and decapsulate its embedding modules by retrieving the buffer.
+    Returns the module with restored embedding (sub) modules.
+    """
+    for child_fqn, child in module.named_children():
+        # perform deserialization on the children first, so that we can replace the child module with
+        # the deserialized module, and then replace it in the parent
+        child = decapsulate_ir_modules(
+            module=child, serializer=serializer, device=device
+        )
+        # replace the child module with deserialized one if applicable
+        setattr(module, child_fqn, child)
+
+    # only deserialize if the module has ir_metadata buffer, otherwise return as is
+    # we use "ir_metadata" as a convention to identify the deserializable module
+    if "ir_metadata" in dict(module.named_buffers()):
+        module = serializer.decapsulate_module(module, device)
+    return module
 
 
 def serialize_embedding_modules(