Update on "[dtensor] enable RMSprop optimizer foreach support"

wanchaol · wanchaol · commit 3ca5917a09cc · 2023-11-20T14:04:38.000-08:00
as titled

cc H-Huang awgu kwen2501 fegin fduwjj wz337 wconstab mrshenli zhaojuanmao rohan-varma kiukchung d4l3k lucasllc XilunWu

[ghstack-poisoned]
diff --git a/torch/distributed/_tensor/ops/pointwise_ops.py b/torch/distributed/_tensor/ops/pointwise_ops.py
@@ -520,23 +520,23 @@ def linear_pointwise_strategy(mesh: DeviceMesh, op_schema: OpSchema) -> Strategy
 
 # TODO: add all for_each ops
 for_each_ops = [
-    aten._foreach_addcmul_.Scalar,
-    aten._foreach_addcmul_.ScalarList,
-    aten._foreach_addcmul_.Tensor,
-    aten._foreach_addcmul.Scalar,
     aten._foreach_addcdiv_.Scalar,
     aten._foreach_addcdiv_.ScalarList,
     aten._foreach_addcdiv_.Tensor,
+    aten._foreach_addcmul.Scalar,
+    aten._foreach_addcmul_.Scalar,
+    aten._foreach_addcmul_.ScalarList,
+    aten._foreach_addcmul_.Tensor,
     aten._foreach_div_.ScalarList,
     aten._foreach_lerp_.Scalar,
     aten._foreach_maximum_.List,
+    aten._foreach_mul.Scalar,
+    aten._foreach_mul.List,
     aten._foreach_mul_.Scalar,
     aten._foreach_mul_.ScalarList,
     aten._foreach_mul_.List,
-    aten._foreach_mul.Scalar,
-    aten._foreach_mul.List,
-    aten._foreach_neg_.default,
     aten._foreach_neg.default,
+    aten._foreach_neg_.default,
     aten._foreach_reciprocal_.default,
     aten._foreach_sub_.Scalar,
     aten._foreach_sqrt.default,
