DTensor: add more foreach ops to supported sharding prop list (#132066)

bdhirsh · pytorchmergebot · commit b465a5843b92 · 2024-08-05T13:51:59.000Z
fixes #132016. Right now if you run an op that DTensor has no sharding prop rule, **and** that op accepts non-trivial pytrees of inputs tensors as arguments, DTensor can end up infinite looping before it has the chance to error due to not having a sharding prop rule. This PR doesn't fix the problem, but adds rules for the culprit ops (missing foreach ops) Pull Request resolved: #132066 Approved by: https://github.com/wanchaol
diff --git a/torch/distributed/_tensor/ops/_pointwise_ops.py b/torch/distributed/_tensor/ops/_pointwise_ops.py
@@ -545,10 +545,18 @@ def linear_pointwise_strategy(mesh: DeviceMesh, op_schema: OpSchema) -> Strategy
     aten._foreach_clamp_max_.Scalar,
     aten._foreach_clamp_min_.Scalar,
     aten._foreach_div_.List,
+    aten._foreach_div_.Scalar,
     aten._foreach_div_.ScalarList,
+    aten._foreach_div_.Tensor,
+    aten._foreach_div.List,
+    aten._foreach_div.Scalar,
+    aten._foreach_div.ScalarList,
+    aten._foreach_div.Tensor,
     aten._foreach_lerp_.Scalar,
     aten._foreach_maximum_.List,
     aten._foreach_mul.Scalar,
+    aten._foreach_mul.ScalarList,
+    aten._foreach_mul.Tensor,
     aten._foreach_mul.List,
     aten._foreach_mul_.Scalar,
     aten._foreach_mul_.ScalarList,
