[dtensor] fix side-effect on dtype for _like ops (pytorch#146869)

fixes pytorch#146749

Pull Request resolved: pytorch#146869
Approved by: https://github.com/yifuwang, https://github.com/janeyx99, https://github.com/ngimel

Author: tianyu-l

test/distributed/tensor/test_tensor_ops.py

@@ -226,9 +226,12 @@ def test_zeros_like(self):
 
         input_tensor = torch.randn(4, 8, requires_grad=True)
         dist_tensor = DTensor.from_local(input_tensor, device_mesh, shard_spec)
-        zeros_like_dt = torch.zeros_like(dist_tensor)
-        zeros_expected = torch.zeros(4, 8)
+        zeros_like_dt = torch.zeros_like(dist_tensor, dtype=torch.bfloat16)
+        zeros_expected = torch.zeros(4, 8, dtype=torch.bfloat16)
         self.assertEqual(zeros_expected, zeros_like_dt.to_local())
+        # make sure there is no side effect on the input tensor dtype
+        self.assertEqual(dist_tensor.dtype, torch.float32)
+        self.assertEqual(zeros_like_dt.dtype, torch.bfloat16)
 
     @with_comms
     @skip_if_lt_x_gpu(4)

torch/distributed/tensor/_ops/_tensor_ops.py

@@ -148,23 +148,16 @@ def create_like_strategy(mesh: DeviceMesh, op_schema: OpSchema) -> StrategyType:
     assert isinstance(select_strategy, OpStrategy)
     for arg_strategy in select_strategy.strategies:
         arg_spec = arg_strategy.output_spec
-        if is_tensor_partial(arg_spec):
-            # if the arg_spec have partial, accept partial
-            # in the input_specs but output replicate for
-            # those corresponding mesh dims
-            output_spec = DTensorSpec(
-                mesh=arg_spec.mesh,
-                placements=tuple(
-                    Replicate() if isinstance(p, Partial) else p
-                    for p in arg_spec.placements
-                ),
-            )
-            create_like_strategy.strategies.append(
-                PlacementStrategy(output_specs=output_spec, input_specs=(arg_spec,))
-            )
-
-        else:
-            create_like_strategy.strategies.append(PlacementStrategy(arg_spec))
+        output_spec = DTensorSpec(
+            mesh=arg_spec.mesh,
+            placements=tuple(
+                Replicate() if isinstance(p, Partial) else p
+                for p in arg_spec.placements
+            ),
+        )
+        create_like_strategy.strategies.append(
+            PlacementStrategy(output_specs=output_spec, input_specs=(arg_spec,))
+        )
 
     return create_like_strategy