in_backward_optimizer_filter on torchrec callsites (#899)

Summary:
Pull Request resolved: #899

X-link: pytorch/torchsnapshot#134

D41964643 , but separate diff to mitigate target determinator

Reviewed By: YLGH

Differential Revision: D42061412

fbshipit-source-id: ea3e0f9aa1739ee718e1dea2664deadb12d44df0

Author: colin2328

contrib/dynamic_embedding/tests/test_integral_precision.py

@@ -1,3 +1,10 @@
+#!/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
@@ -14,6 +21,7 @@
 
 from torchrec.distributed.planner import EmbeddingShardingPlanner, Topology
 from torchrec.optim.keyed import CombinedOptimizer, KeyedOptimizerWrapper
+from torchrec.optim.optimizers import in_backward_optimizer_filter
 
 from torchrec_dynamic_embedding.id_transformer_group import IDTransformerGroup
 from utils import init_dist, register_memory_io
@@ -93,7 +101,7 @@ def get_dmp(model):
             model = DMP(module=model, device=device, plan=plan, sharders=sharders)
 
             dense_optimizer = KeyedOptimizerWrapper(
-                dict(model.named_parameters()),
+                dict(in_backward_optimizer_filter(model.named_parameters())),
                 lambda params: torch.optim.Adam(params, lr=1e-1),
             )
             optimizer = CombinedOptimizer([model.fused_optimizer, dense_optimizer])

examples/bert4rec/bert4rec_main.py

@@ -26,6 +26,7 @@
 from torchrec.distributed.model_parallel import DistributedModelParallel as DMP
 from torchrec.distributed.types import ModuleSharder
 from torchrec.optim.keyed import CombinedOptimizer, KeyedOptimizerWrapper
+from torchrec.optim.optimizers import in_backward_optimizer_filter
 from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
 from tqdm import tqdm
 
@@ -497,11 +498,12 @@ def main(argv: List[str]) -> None:
             ],
         )
         dense_optimizer = KeyedOptimizerWrapper(
-            dict(model.named_parameters()),
+            dict(in_backward_optimizer_filter(model.named_parameters())),
             lambda params: optim.Adam(
                 params, lr=args.lr, weight_decay=args.weight_decay
             ),
         )
+
         optimizer = CombinedOptimizer([model.fused_optimizer, dense_optimizer])
     else:
         device_ids = [rank] if backend == "nccl" else None

examples/golden_training/train_dlrm.py

@@ -27,6 +27,7 @@
 from torchrec.modules.embedding_modules import EmbeddingBagCollection
 from torchrec.optim.apply_optimizer_in_backward import apply_optimizer_in_backward
 from torchrec.optim.keyed import KeyedOptimizerWrapper
+from torchrec.optim.optimizers import in_backward_optimizer_filter
 from torchrec.optim.rowwise_adagrad import RowWiseAdagrad
 from tqdm import tqdm
 
@@ -132,7 +133,7 @@ def train(
     )
 
     non_fused_optimizer = KeyedOptimizerWrapper(
-        dict(model.named_parameters()),
+        dict(in_backward_optimizer_filter(model.named_parameters())),
         lambda params: torch.optim.Adagrad(params, lr=learning_rate),
     )
     # Overlap comm/compute/device transfer during training through train_pipeline

examples/nvt_dataloader/train_torchrec.py

@@ -19,8 +19,6 @@
 import torchrec
 import torchrec.distributed as trec_dist
 import torchrec.optim as trec_optim
-
-from fbgemm_gpu.split_embedding_configs import EmbOptimType as OptimType
 from nvt_binary_dataloader import NvtBinaryDataloader
 from pyre_extensions import none_throws
 from torchrec import EmbeddingBagCollection
@@ -40,6 +38,7 @@
 from torchrec.modules.embedding_configs import EmbeddingBagConfig
 from torchrec.modules.fused_embedding_modules import fuse_embedding_optimizer
 from torchrec.optim.keyed import KeyedOptimizerWrapper
+from torchrec.optim.optimizers import in_backward_optimizer_filter
 
 
 def parse_args(argv: List[str]) -> argparse.Namespace:
@@ -270,7 +269,7 @@ def main(argv: List[str]):
     )
 
     non_fused_optimizer = KeyedOptimizerWrapper(
-        dict(model.named_parameters()),
+        dict(in_backward_optimizer_filter(model.named_parameters())),
         lambda params: torch.optim.Adagrad(params, lr=args.learning_rate)
         if args.adagrad
         else torch.optim.SGD(params, lr=args.learning_rate),

examples/ray/train_torchrec.py

@@ -24,6 +24,7 @@
 from torchrec.modules.embedding_configs import EmbeddingBagConfig
 from torchrec.modules.embedding_modules import EmbeddingBagCollection
 from torchrec.optim.keyed import KeyedOptimizerWrapper
+from torchrec.optim.optimizers import in_backward_optimizer_filter
 from tqdm import tqdm
 
 
@@ -110,7 +111,7 @@ def train(
 
     # Overlap comm/compute/device transfer during training through train_pipeline
     non_fused_optimizer = KeyedOptimizerWrapper(
-        dict(model.named_parameters()),
+        dict(in_backward_optimizer_filter(model.named_parameters())),
         lambda params: torch.optim.Adagrad(params, lr=learning_rate),
     )
     train_pipeline = TrainPipelineSparseDist(

examples/torcharrow/run.py

@@ -20,6 +20,7 @@
 from torchrec.models.dlrm import DLRM
 from torchrec.modules.embedding_configs import EmbeddingBagConfig
 from torchrec.optim.keyed import KeyedOptimizerWrapper
+from torchrec.optim.optimizers import in_backward_optimizer_filter
 
 
 @record
@@ -94,7 +95,7 @@ def main(
     )
 
     optimizer = KeyedOptimizerWrapper(
-        dict(model.named_parameters()),
+        dict(in_backward_optimizer_filter(model.named_parameters())),
         lambda params: torch.optim.SGD(params, lr=0.01),
     )
 

test_installation.py

@@ -18,6 +18,7 @@
 from torchrec.models.dlrm import DLRM
 from torchrec.modules.embedding_configs import EmbeddingBagConfig
 from torchrec.optim.keyed import KeyedOptimizerWrapper
+from torchrec.optim.optimizers import in_backward_optimizer_filter
 
 if sys.platform not in ["linux", "linux2"]:
     raise EnvironmentError(
@@ -118,7 +119,7 @@ def main(argv: List[str]) -> None:
         device=device,
     )
     optimizer = KeyedOptimizerWrapper(
-        dict(model.named_parameters()),
+        dict(in_backward_optimizer_filter(model.named_parameters())),
         lambda params: torch.optim.SGD(params, lr=0.01),
     )
 

torchrec/distributed/test_utils/test_sharding.py

@@ -43,6 +43,7 @@
 from torchrec.modules.embedding_configs import BaseEmbeddingConfig, EmbeddingBagConfig
 from torchrec.optim.apply_optimizer_in_backward import apply_optimizer_in_backward
 from torchrec.optim.keyed import CombinedOptimizer, KeyedOptimizerWrapper
+from torchrec.optim.optimizers import in_backward_optimizer_filter
 from typing_extensions import Protocol
 
 
@@ -343,7 +344,7 @@ def sharding_single_rank_test(
         )
 
         dense_optim = KeyedOptimizerWrapper(
-            dict(local_model.named_parameters()),
+            dict(in_backward_optimizer_filter(local_model.named_parameters())),
             lambda params: torch.optim.SGD(params, lr=0.1),
         )
         local_opt = CombinedOptimizer([local_model.fused_optimizer, dense_optim])

torchrec/distributed/tests/test_train_pipeline.py

@@ -44,6 +44,7 @@
 from torchrec.modules.embedding_modules import EmbeddingBagCollection
 
 from torchrec.optim.keyed import KeyedOptimizerWrapper
+from torchrec.optim.optimizers import in_backward_optimizer_filter
 from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
 from torchrec.streamable import Pipelineable
 from torchrec.test_utils import get_free_port, init_distributed_single_host
@@ -197,7 +198,7 @@ def _test_feature_processor_helper(
         copy_state_dict(unsharded_model.state_dict(), distributed_model.state_dict())
         optimizer_cpu = optim.SGD(unsharded_model.parameters(), lr=0.1)
         optimizer_distributed = KeyedOptimizerWrapper(
-            dict(distributed_model.named_parameters()),
+            dict(in_backward_optimizer_filter(distributed_model.named_parameters())),
             lambda params: optim.SGD(params, lr=0.1),
         )
         pipeline = TrainPipelineSparseDist(
@@ -289,7 +290,7 @@ def _test_move_cpu_gpu_helper(
         )
         optimizer_cpu = optim.SGD(model_cpu.parameters(), lr=0.1)
         optimizer_distributed = KeyedOptimizerWrapper(
-            dict(distributed_model.named_parameters()),
+            dict(in_backward_optimizer_filter(distributed_model.named_parameters())),
             lambda params: optim.SGD(params, lr=0.1),
         )
         pipeline = TrainPipelineSparseDist(