[dtensor] fix foreach_norm when ord is 2 (pytorch#130753)

as titled, fixed a case when passing ord as 2 (default value), the op
dispatching does not receive the default value case

We simply check if the args schema receiving a `ord` field or not

Pull Request resolved: pytorch#130753
Approved by: https://github.com/awgu

Author: wanchaol

test/distributed/_tensor/test_math_ops.py

@@ -437,6 +437,25 @@ def test_shard0_svd(self):
         self.assertEqual(len(comm_counts), 1)
         self.assertEqual(comm_counts[funcol.all_gather_into_tensor], 1)
 
+    @with_comms
+    def test_foreach_norm(self):
+        device_mesh = self.build_device_mesh()
+
+        grad0 = torch.randn(12, 8)
+        grad1 = torch.randn(8, 8)
+
+        sharded_grad0 = distribute_tensor(grad0, device_mesh, [Shard(0)])
+        sharded_grad1 = distribute_tensor(grad1, device_mesh, [Shard(0)])
+
+        # non-sharded op
+        out = torch.ops.aten._foreach_norm([grad0, grad1], 2)
+
+        # sharded op
+        sharded_out = torch.ops.aten._foreach_norm([sharded_grad0, sharded_grad1], 2)
+
+        for o, so in zip(out, sharded_out):
+            self.assertEqual(so.full_tensor(), o)
+
 
 if __name__ == "__main__":
     run_tests()

torch/distributed/_tensor/ops/math_ops.py

@@ -397,7 +397,7 @@ def foreach_norm_strategy(mesh: DeviceMesh, op_schema: OpSchema) -> TupleStrateg
     args_schema = op_schema.args_schema
     input_tuple_strategy = args_schema[0]
     assert isinstance(input_tuple_strategy, TupleStrategy)
-    norm_type = args_schema[1]
+    norm_type = args_schema[1] if len(args_schema) > 1 else 2
     assert isinstance(norm_type, (int, float, str)), f"{norm_type}"
     output_tuple_strategy_childs: List[OpStrategy] = []
     for op_strategy in input_tuple_strategy.childs: