Introduce Serializer in TorchRec for torch.export PEA (#1848)

Summary:
Pull Request resolved: #1848

Adding a serializer class to support thrift serialization of PEA configs for regenerating the eager module after torch.export through saving the metadata as a bytes buffer in the module.

Reviewed By: dstaay-fb

Differential Revision: D55661704

fbshipit-source-id: 6ea47e495eaad032f17ad6c5f45206259c8009b3

Author: PaulZhang12

torchrec/ir/types.py

@@ -0,0 +1,42 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+#!/usr/bin/env python3
+
+import abc
+from typing import Any, Dict, Type
+
+from torch import nn
+
+
+class SerializerInterface(abc.ABC):
+    """
+    Interface for Serializer classes for torch.export IR.
+    """
+
+    @classmethod
+    @property
+    # pyre-ignore [3]: Returning `None` but type `Any` is specified.
+    def module_to_serializer_cls(cls) -> Dict[Type[nn.Module], Type[Any]]:
+        raise NotImplementedError
+
+    @classmethod
+    @abc.abstractmethod
+    # pyre-ignore [3]: Returning `None` but type `Any` is specified.
+    def serialize(
+        cls,
+        module: nn.Module,
+    ) -> Any:
+        # Take the eager embedding module and generate bytes in buffer
+        pass
+
+    @classmethod
+    @abc.abstractmethod
+    # pyre-ignore [2]: Parameter `input` must have a type other than `Any`.
+    def deserialize(cls, input: Any) -> nn.Module:
+        # Take the bytes in the buffer and regenerate the eager embedding module
+        pass

torchrec/ir/utils.py

@@ -0,0 +1,46 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+#!/usr/bin/env python3
+
+from typing import Type
+
+from torch import nn
+from torchrec.ir.types import SerializerInterface
+
+
+# TODO: Replace the default interface with the python dataclass interface
+DEFAULT_SERIALIZER_CLS = SerializerInterface
+
+
+def serialize_embedding_modules(
+    model: nn.Module,
+    serializer_cls: Type[SerializerInterface] = DEFAULT_SERIALIZER_CLS,
+) -> nn.Module:
+    for _, module in model.named_modules():
+        if type(module) in serializer_cls.module_to_serializer_cls:
+            serialized_module = serializer_cls.serialize(module)
+            module.register_buffer("ir_metadata", serialized_module, persistent=False)
+
+    return model
+
+
+def deserialize_embedding_modules(
+    model: nn.Module,
+    serializer_cls: Type[SerializerInterface] = DEFAULT_SERIALIZER_CLS,
+) -> nn.Module:
+    fqn_to_new_module = {}
+    for name, module in model.named_modules():
+        if "ir_metadata" in dict(module.named_buffers()):
+            serialized_module = dict(module.named_buffers())["ir_metadata"]
+            deserialized_module = serializer_cls.deserialize(serialized_module)
+            fqn_to_new_module[name] = deserialized_module
+
+    for fqn, new_module in fqn_to_new_module.items():
+        setattr(model, fqn, new_module)
+
+    return model