Fix FeatureProcessor device only to meta

Summary:
D56021085 - This fixed copying FP parameters to meta device when sharding model on meta
Weird thing is though that FP parameters are not sparse parameters, they are dense. Therefore, they shouldn’t be moved to meta device as a result of sharding.
https://fburl.com/code/xuv9s5k2 - AIMP assumes only sparse params are on meta device.

However, the FP parameters should be properly moved when using .to(). That is not the current case, as FeatureProcessorCollections use ParameterDict but also has a mirroring Dict[str, nn.Parameter] for bypassing TorchScript issues with ParameterDict. Therefore, when .to() is used on a model, only the registered ParameterDict will bem oved but not the mirroring Dict. This diff overrides nn.Module _apply to make sure the ParameterDict and Dict are in sync.

Reviewed By: IvanKobzarev, gnahzg

Differential Revision: D56492970

fbshipit-source-id: 08e437ae618cab67c2dceb2d8ccb7547132eafbe

torchrec/distributed/quant_embeddingbag.py

@@ -411,7 +411,7 @@ def __init__(
                 # Generic copy, for example initailized on cpu but -> sharding as meta
                 self.feature_processors_per_rank.append(
                     copy.deepcopy(feature_processor)
-                    if device_type == feature_processor_device
+                    if device_type == "meta"
                     else copy_to_device(
                         feature_processor,
                         feature_processor_device,

torchrec/distributed/tests/test_infer_shardings.py

@@ -2019,6 +2019,8 @@ def test_sharded_quant_fp_ebc_tw_meta(self) -> None:
             if isinstance(input, torch.Tensor):
                 inputs[i] = input.to(torch.device("meta"))
 
+        # move dense params also to meta
+        sharded_model.to("meta")
         sharded_model(*inputs)
         # Don't care about the output since we are sharding on meta
 

torchrec/distributed/tests/test_pt2.py

@@ -322,6 +322,9 @@ def test_sharded_quant_fpebc_non_strict_export(self) -> None:
         )
         kjt = input_kjts[0]
         kjt = kjt.to("meta")
+        # Move FP parameters
+        sharded_model.to("meta")
+
         sharded_model(kjt.values(), kjt.lengths())
 
         ep = torch.export.export(

torchrec/modules/feature_processor_.py

@@ -201,3 +201,12 @@ def copy(self, device: torch.device) -> nn.Module:
             self.position_weights_dict[key] = self.position_weights[key]
 
         return self
+
+    # Override to make sure position_weights and position_weights_dict are in sync
+    # pyre-ignore [2]
+    def _apply(self, *args, **kwargs) -> nn.Module:
+        super()._apply(*args, **kwargs)
+        for k, param in self.position_weights.items():
+            self.position_weights_dict[k] = param
+
+        return self

torchrec/modules/tests/test_feature_processor_.py

@@ -182,3 +182,26 @@ def test_copy(self) -> None:
         self.assertTrue(
             all(param.is_meta for param in res.position_weights_dict.values())
         )
+
+    def test_to(self) -> None:
+        pwmc = PositionWeightedModuleCollection(
+            max_feature_lengths={"f1": 10, "f2": 10},
+            device=torch.device("cpu"),
+        )
+
+        self.assertTrue(
+            all(param.device.type == "cpu" for param in pwmc.position_weights.values())
+        )
+        self.assertTrue(
+            all(
+                param.device.type == "cpu"
+                for param in pwmc.position_weights_dict.values()
+            )
+        )
+
+        pwmc.to("meta")
+
+        self.assertTrue(all(param.is_meta for param in pwmc.position_weights.values()))
+        self.assertTrue(
+            all(param.is_meta for param in pwmc.position_weights_dict.values())
+        )