[Semi-Auto] Improve Global Clip (#65928)

JZ-LIANG · winter-wang · web-flow · commit afe4d0b726d0 · 2024-07-15T11:19:13.000+08:00
* add code

* bugfix

---------

Co-authored-by: winter-wang &lt;1030748926@qq.com&gt;
diff --git a/paddle/phi/infermeta/spmd_rules/rules.h b/paddle/phi/infermeta/spmd_rules/rules.h
@@ -47,6 +47,7 @@ limitations under the License. */
 #include "paddle/phi/infermeta/spmd_rules/slice.h"
 #include "paddle/phi/infermeta/spmd_rules/softmax.h"
 #include "paddle/phi/infermeta/spmd_rules/split.h"
+#include "paddle/phi/infermeta/spmd_rules/squared_l2_norm.h"
 #include "paddle/phi/infermeta/spmd_rules/squeeze.h"
 #include "paddle/phi/infermeta/spmd_rules/stack.h"
 #include "paddle/phi/infermeta/spmd_rules/tile.h"
diff --git a/paddle/phi/infermeta/spmd_rules/squared_l2_norm.cc b/paddle/phi/infermeta/spmd_rules/squared_l2_norm.cc
@@ -0,0 +1,44 @@
+/* Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include "paddle/phi/infermeta/spmd_rules/squared_l2_norm.h"
+#include "paddle/phi/infermeta/spmd_rules/reduction.h"
+
+#include "glog/logging.h"
+
+#include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
+#include "paddle/phi/core/distributed/auto_parallel/inferspmd_utils.h"
+#include "paddle/phi/core/distributed/auto_parallel/utils.h"
+#include "paddle/phi/infermeta/spmd_rules/utils.h"
+
+namespace phi {
+namespace distributed {
+
+using phi::distributed::auto_parallel::str_join;
+
+SpmdInfo SquaredL2NormInferSpmd(const DistMetaTensor& x) {
+  VLOG(4) << "SquaredL2NormInferSpmd:";
+  VLOG(4) << "Using ReductionInferSpmd Rule as interal implement.";
+  SpmdInfo info = ReductionInferSpmdBase(
+      x, {}, false, static_cast<int>(ReduceType::kRedSum));
+  // NOTE: reduce output is 0D tensor which has a dims_mapping as {}, while
+  // output of squared_l2_norm is a tensor with shape [1] therefore it need to
+  // have a dims_mapping as {-1}.
+  auto& out_dist_dst = PADDLE_GET(TensorDistAttr, info.second[0]);
+  out_dist_dst.set_dims_mapping({-1});
+  return info;
+}
+
+}  // namespace distributed
+}  // namespace phi
diff --git a/paddle/phi/infermeta/spmd_rules/squared_l2_norm.h b/paddle/phi/infermeta/spmd_rules/squared_l2_norm.h
@@ -0,0 +1,29 @@
+/* Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#pragma once
+
+#include <vector>
+
+#include "paddle/phi/common/int_array.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_meta_tensor.h"
+#include "paddle/phi/core/distributed/type_defs.h"
+
+namespace phi {
+namespace distributed {
+
+SpmdInfo SquaredL2NormInferSpmd(const DistMetaTensor& x);
+
+}  // namespace distributed
+}  // namespace phi
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -4260,6 +4260,7 @@
   output : Tensor(out)
   infer_meta :
     func : SquaredL2NormInferMeta
+    spmd_rule : SquaredL2NormInferSpmd
   kernel :
     func : squared_l2_norm
   backward : squared_l2_norm_grad
