add in_backward_optimizer_filter to work with in_backward_optimizers (#315)

Summary:
X-link: facebookresearch/dlrm#315

Pull Request resolved: #892

as title

Reviewed By: lequytra, YLGH

Differential Revision: D42009102

fbshipit-source-id: a51cbe59a66a1ce4d82f5c46d39eca89b326be7a

Author: colin2328

torchrec/optim/optimizers.py

@@ -7,13 +7,32 @@
 
 #!/usr/bin/env python3
 
-from typing import Iterable
+from typing import Iterable, Iterator, Tuple
 
 import torch
+from torch import nn
 
 from torch.optim.optimizer import Optimizer
 
 
+def in_backward_optimizer_filter(
+    named_parameters: Iterator[Tuple[str, nn.Parameter]], include: bool = False
+) -> Iterator[Tuple[str, nn.Parameter]]:
+    """
+    Filters named_parameters for whether they are or or not params that use
+    the in_backward_optimizer.
+
+    Args:
+    named_parameters(Iterator[Tuple[str, nn.Parameter]]): named_parameters
+    include(bool): If true, only yields params with in_backward_optimizer. If false, returns the outside set
+        Defaults to include params that are not in_backward (False)
+    """
+    for fqn, param in named_parameters:
+        # TODO: change to _in_backward_optimizer
+        if hasattr(param, "_overlapped_optimizer") == include:
+            yield fqn, param
+
+
 class SGD(Optimizer):
     r"""
     Placeholder for SGD. This optimizer will not functionally run.

torchrec/optim/tests/test_optim.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.
+
+import unittest
+
+import torch
+
+from torchrec.modules.embedding_configs import EmbeddingBagConfig
+from torchrec.modules.embedding_modules import EmbeddingBagCollection
+from torchrec.optim.apply_optimizer_in_backward import apply_optimizer_in_backward
+from torchrec.optim.optimizers import in_backward_optimizer_filter
+
+
+class TestInBackwardOptimizerFilter(unittest.TestCase):
+    def test_in_backward_optimizer_filter(self) -> None:
+        ebc = EmbeddingBagCollection(
+            tables=[
+                EmbeddingBagConfig(
+                    name="t1", embedding_dim=4, num_embeddings=2, feature_names=["f1"]
+                ),
+                EmbeddingBagConfig(
+                    name="t2", embedding_dim=4, num_embeddings=2, feature_names=["f2"]
+                ),
+            ]
+        )
+        apply_optimizer_in_backward(
+            torch.optim.SGD,
+            ebc.embedding_bags["t1"].parameters(),
+            optimizer_kwargs={"lr": 1.0},
+        )
+        in_backward_params = dict(
+            in_backward_optimizer_filter(ebc.named_parameters(), include=True)
+        )
+        non_in_backward_params = dict(
+            in_backward_optimizer_filter(ebc.named_parameters(), include=False)
+        )
+        assert set(in_backward_params.keys()) == {"embedding_bags.t1.weight"}
+        assert set(non_in_backward_params.keys()) == {"embedding_bags.t2.weight"}