PaddlePaddle
diff --git a/‎paddle/fluid/imperative/gradient_accumulator.cc
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion b/‎paddle/fluid/imperative/gradient_accumulator.cc
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion
diff --git a/‎paddle/fluid/pybind/auto_parallel_py.cc
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions b/‎paddle/fluid/pybind/auto_parallel_py.cc
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions
diff --git a/‎paddle/fluid/pybind/eager_functions.cc
Copy file name to clipboard
+3-2Lines changed: 3 additions & 2 deletions b/‎paddle/fluid/pybind/eager_functions.cc
Copy file name to clipboard
+3-2Lines changed: 3 additions & 2 deletions
diff --git a/‎paddle/phi/core/distributed/auto_parallel/process_mesh.h
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions b/‎paddle/phi/core/distributed/auto_parallel/process_mesh.h
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions
diff --git a/‎paddle/phi/infermeta/spmd_rules/cast.cc
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion b/‎paddle/phi/infermeta/spmd_rules/cast.cc
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion
diff --git a/‎paddle/phi/infermeta/spmd_rules/elementwise.cc
Copy file name to clipboard
+49Lines changed: 49 additions & 0 deletions b/‎paddle/phi/infermeta/spmd_rules/elementwise.cc
Copy file name to clipboard
+49Lines changed: 49 additions & 0 deletions
diff --git a/‎paddle/phi/infermeta/spmd_rules/elementwise.h
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions b/‎paddle/phi/infermeta/spmd_rules/elementwise.h
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions
diff --git a/‎python/paddle/amp/auto_cast.py
Copy file name to clipboard
+27-1Lines changed: 27 additions & 1 deletion b/‎python/paddle/amp/auto_cast.py
Copy file name to clipboard
+27-1Lines changed: 27 additions & 1 deletion
diff --git a/‎python/paddle/amp/grad_scaler.py
Copy file name to clipboard
+2-1Lines changed: 2 additions & 1 deletion b/‎python/paddle/amp/grad_scaler.py
Copy file name to clipboard
+2-1Lines changed: 2 additions & 1 deletion
diff --git a/‎python/paddle/distributed/__init__.py
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions b/‎python/paddle/distributed/__init__.py
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions
@@ -204,7 +204,7 @@ void TensorAdd(const VarType& src, VarType* dst) {
   if (data_type == framework::DataTypeTrait<T>::DataType()) {                  \
     auto cpu_ctx = static_cast<CONTEXT*>(                                      \
         platform::DeviceContextPool::Instance().Get(place));                   \
-    phi::AddKernel<T, CONTEXT>(*cpu_ctx, src_tensor, *dst_tensor, dst_tensor); \
+    phi::AddKernel<T, CONTEXT>(*cpu_ctx, *dst_tensor, src_tensor, dst_tensor); \
     return;                                                                    \
   }
 
 
@@ -278,6 +278,7 @@ void BindAutoParallel(py::module *m) {
                 return ProcessMesh(self);
               },
               py::arg("memo"))
+          .def("__hash__", &ProcessMesh::hash)
           .def("__str__", &ProcessMesh::to_string);
 
   g_process_mesh_pytype = reinterpret_cast<PyTypeObject *>(process_mesh.ptr());
 
@@ -1338,8 +1338,9 @@ static PyObject* eager_api_set_master_grads(PyObject* self,
                           "Detected nullptr grad"
                           "Please check if you have manually cleared"
                           "the grad inside autograd_meta"));
-    if ((*grad).initialized() && ((*grad).dtype() == phi::DataType::FLOAT16 ||
-                                  (*grad).dtype() == phi::DataType::BFLOAT16)) {
+    if (((*grad).initialized() || (*grad).is_dist_tensor()) &&
+        ((*grad).dtype() == phi::DataType::FLOAT16 ||
+         (*grad).dtype() == phi::DataType::BFLOAT16)) {
       auto master_grad =
           paddle::experimental::cast(*grad, phi::DataType::FLOAT32);
       grad->set_impl(master_grad.impl());
 
@@ -67,6 +67,8 @@ class ProcessMesh {
   bool empty() const { return (shape_.empty() || process_ids_.empty()); }
   bool contains(int64_t process_id) const;
 
+  size_t hash() const { return std::hash<std::string>{}(to_string()); }
+
   // ProcessMesh from_string(const std::string& mesh_str);
   std::string to_string() const;
 
 
@@ -15,7 +15,7 @@ limitations under the License. */
 namespace phi {
 namespace distributed {
 SpmdInfo CastInferSpmd(const DistMetaTensor& x, phi::DataType dtype) {
-  return ElementwiseUnaryInferSpmd(x);
+  return ElementwiseUnaryWithPartialInferSpmd(x);
 }
 }  // namespace distributed
 }  // namespace phi
@@ -124,6 +124,55 @@ SpmdInfo ElementwiseUnaryInferSpmd(const DistMetaTensor& x) {
   return {{x_dst_dist_attr}, {out_dist_attr}};
 }
 
+// NOTE(lizhiyu): This function is only for `cast` right now to support partial
+// propagatation
+SpmdInfo ElementwiseUnaryWithPartialInferSpmd(const DistMetaTensor& x) {
+  // Step0: Verify Input Args Based on Elementwise Logic
+  auto x_shape = common::vectorize(x.dims());
+  int x_ndim = x_shape.size();
+  TensorDistAttr x_dist_attr_src = x.dist_attr();
+  std::vector<int64_t> x_dims_mapping = x_dist_attr_src.dims_mapping();
+  PADDLE_ENFORCE_EQ(x_ndim,
+                    x_dims_mapping.size(),
+                    phi::errors::InvalidArgument(
+                        "ElementwiseUnary, The Tensor X's rank [%d] and X's "
+                        "dims_mapping size [%d] are not matched.",
+                        x_ndim,
+                        x_dims_mapping.size()));
+
+  // Step1: Build Einsum Notation
+  std::string alphabet = "abcdefghijklmnopqrstuvwxyz";
+  std::string x_axes = GetBroadcastAxes(x_ndim, x_ndim, alphabet);
+  std::string out_axes = x_axes;
+
+  // Step2: Sharding Propogation
+  // Step2.1: Merge input shardings
+  std::pair<std::string, std::vector<int64_t>> axes_sharding_info(
+      x_axes, x_dims_mapping);
+  std::unordered_map<std::string, int64_t> axis_to_dim_map =
+      ShardingMergeForTensors({axes_sharding_info});
+
+  // step2.2: Infer output dims mapping from merged input dims mapping
+  std::vector<int64_t> out_dims_mapping =
+      GetDimsMappingForAxes(out_axes, axis_to_dim_map);
+
+  // initialize output dist_attr's process_mesh, batch_dim and dynamic dims with
+  // input dist_attr.
+  TensorDistAttr out_dist_attr = CopyTensorDistAttrForOutput(x_dist_attr_src);
+  out_dist_attr.set_dims_mapping(out_dims_mapping);
+  out_dist_attr.set_partial_status(x_dist_attr_src.partial_status());
+  TensorDistAttr x_dst_dist_attr = CopyTensorDistAttrForOutput(x_dist_attr_src);
+  x_dst_dist_attr.set_dims_mapping(out_dims_mapping);
+  x_dst_dist_attr.set_partial_status(x_dist_attr_src.partial_status());
+
+  VLOG(4) << "ElementwiseWithPartialSPMDRule InferForward:";
+  VLOG(4) << "Input0 shape: [" << str_join(x_shape) << "] "
+          << "src_dims_mapping: [" << str_join(x_dims_mapping) << "] ";
+  VLOG(4) << "Output dims_mapping: [" + str_join(out_dims_mapping) + "]\n\n";
+
+  return {{x_dst_dist_attr}, {out_dist_attr}};
+}
+
 SpmdInfo ElementwiseUnaryInferSpmdReverse(const DistMetaTensor& x,
                                           const DistMetaTensor& out) {
   // Step0: Verify Input Args Based on Elementwise Logic
 
@@ -24,6 +24,8 @@ namespace distributed {
 
 SpmdInfo ElementwiseUnaryInferSpmd(const DistMetaTensor& x);
 
+SpmdInfo ElementwiseUnaryWithPartialInferSpmd(const DistMetaTensor& x);
+
 SpmdInfo ElementwiseUnaryInferSpmdReverse(const DistMetaTensor& x,
                                           const DistMetaTensor& out);
 
 
@@ -48,6 +48,8 @@ def __init__(self):
         self.model_parameters = []
         self.use_master_grad = False
         self.already_register_final_backward_hook = False
+        self.already_classify_params_meshs = False  # For dist
+        self.mesh2params = {}  # For dist
         self.amp_dtype = 'float32'
 
     def __setattr__(self, name, val):
@@ -423,7 +425,31 @@ def amp_guard(
     ):
 
         def master_grad_hook():
-            core.eager.set_master_grads(amp_global_state().model_parameters)
+            # NOTE(lizhiyu): To support semi-auto of dygraph mode, we must
+            # classify the params of model into different calsses according to their process_mesh.
+            # Otherwise, fault will occur.
+            if not amp_global_state().already_classify_params_meshs:
+                for param in amp_global_state().model_parameters:
+                    if param is not None and param.process_mesh is not None:
+                        if (
+                            param.process_mesh
+                            not in amp_global_state().mesh2params
+                        ):
+                            amp_global_state().mesh2params[
+                                param.process_mesh
+                            ] = [param]
+                        else:
+                            amp_global_state().mesh2params[
+                                param.process_mesh
+                            ].append(param)
+                amp_global_state().already_classify_params_meshs = True
+
+            if len(amp_global_state().mesh2params):
+                for _, params in amp_global_state().mesh2params.items():
+                    core.eager.set_master_grads(params)
+            else:
+                core.eager.set_master_grads(amp_global_state().model_parameters)
+
             amp_global_state().already_register_final_backward_hook = False
 
         core.eager._add_backward_final_hook(master_grad_hook)
 
@@ -193,7 +193,8 @@ def scale(self, var):
                 % (amp_global_state().amp_dtype)
             )
 
-        if not self._enable:
+        # NOTE(lizhiyu): We hack here to avoid changing the `dist_attr` of `self._scale` of 'no-calculation-rank'
+        if not self._enable or not var._is_initialized():
             return var
 
         return var * self._scale
 
@@ -85,6 +85,7 @@
     shard_dataloader,
     shard_layer,
     shard_optimizer,
+    shard_scaler,
     to_static,
     Strategy,
     DistModel,
@@ -169,6 +170,7 @@
     "save_state_dict",
     "load_state_dict",
     "shard_optimizer",
+    "shard_scaler",
     "to_static",
     "Strategy",
     "DistModel",
Original file line number	Diff line number	Diff line change
`@@ -204,7 +204,7 @@ void TensorAdd(const VarType& src, VarType* dst) {`
`204`	`204`	`if (data_type == framework::DataTypeTrait<T>::DataType()) { \`
`205`	`205`	`auto cpu_ctx = static_cast<CONTEXT*>( \`
`206`	`206`	`platform::DeviceContextPool::Instance().Get(place)); \`
`207`		`- phi::AddKernel<T, CONTEXT>(cpu_ctx, src_tensor, dst_tensor, dst_tensor); \`
	`207`	`+ phi::AddKernel<T, CONTEXT>(cpu_ctx, dst_tensor, src_tensor, dst_tensor); \`
`208`	`208`	`return; \`
`209`	`209`	`}`
`210`	`210`
Original file line number	Diff line number	Diff line change
`@@ -15,7 +15,7 @@ limitations under the License. */`
`15`	`15`	`namespace phi {`
`16`	`16`	`namespace distributed {`
`17`	`17`	`SpmdInfo CastInferSpmd(const DistMetaTensor& x, phi::DataType dtype) {`
`18`		`- return ElementwiseUnaryInferSpmd(x);`
	`18`	`+ return ElementwiseUnaryWithPartialInferSpmd(x);`
`19`	`19`	`}`
`20`	`20`	`} // namespace distributed`
`21`	`21`	`} // namespace phi`
Original file line number	Diff line number	Diff line change
`@@ -193,7 +193,8 @@ def scale(self, var):`
`193`	`193`	`% (amp_global_state().amp_dtype)`
`194`	`194`	`)`
`195`	`195`
`196`		`- if not self._enable:`
	`196`	+ # NOTE(lizhiyu): We hack here to avoid changing the `dist_attr` of `self._scale` of 'no-calculation-rank'
	`197`	`+ if not self._enable or not var._is_initialized():`
`197`	`198`	`return var`
`198`	`199`
`199`	`200`	`return var * self._scale`