Fit PIR amp O2 for auto_parallel (#66418)

zhiqiu · web-flow · commit af414f578620 · 2024-07-24T16:07:42.000+08:00
* fix ut

* fix skip_update on global mesh

* refine SetInplaceOutputCorrectDistAttr
diff --git a/paddle/phi/api/lib/data_transform.cc b/paddle/phi/api/lib/data_transform.cc
@@ -809,29 +809,22 @@ void SetInplaceOutputCorrectDistAttr(
     phi::distributed::DistTensor* dist_tensor =
         static_cast<phi::distributed::DistTensor*>(tensor_in.get());
     if (dist_tensor->initialized()) {
-      if (ReshardIsNeededWithPartial(dist_tensor->dist_attr(), dist_attr)) {
-        if (use_general_spmd_rule) {
+      if (use_general_spmd_rule) {
+        if (ReshardIsNeeded(dist_tensor->dist_attr(), dist_attr)) {
           VLOG(6) << "SetInplaceOutputCorrectDistAttr Reshard inplace output"
                   << " to origin dist_attr "
                   << ReshardDebugInfo(*dist_tensor, dist_attr);
           auto* func = phi::distributed::ChooseProperReshardFunction(
               *dist_tensor, dist_attr);
           func->Eval(dev_ctx, *dist_tensor, dist_attr, dist_tensor);
-        } else {
-          // just set correct SPMD dist_attrs
-          VLOG(6) << "SetInplaceOutputCorrectDistAttr input " << tensor.name()
-                  << " set its dist_attr from " << dist_tensor->dist_attr()
-                  << " to " << dist_attr;
-          dist_tensor->unsafe_set_dist_attr(dist_attr);
+          return;
         }
       }
-    } else {
-      VLOG(6) << "SetInplaceOutputCorrectDistAttr has"
-              << " uninitialized DistTensor input " << tensor.name()
-              << ", just set its dist_attr from " << dist_tensor->dist_attr()
-              << " to " << dist_attr;
-      dist_tensor->unsafe_set_dist_attr(dist_attr);
     }
+    VLOG(6) << "SetInplaceOutputCorrectDistAttr for tensor " << tensor.name()
+            << ", just set its dist_attr from " << dist_tensor->dist_attr()
+            << " to " << dist_attr;
+    dist_tensor->unsafe_set_dist_attr(dist_attr);
   }
 }
 
@@ -859,30 +852,23 @@ void SetInplaceOutputCorrectDistAttr(
       phi::distributed::DistTensor* dist_tensor =
           static_cast<phi::distributed::DistTensor*>(tensor_in.get());
       if (dist_tensor->initialized()) {
-        if (ReshardIsNeededWithPartial(dist_tensor->dist_attr(),
-                                       dist_attr[i])) {
-          if (use_general_spmd_rule) {
+        if (use_general_spmd_rule) {
+          if (ReshardIsNeededWithPartial(dist_tensor->dist_attr(),
+                                         dist_attr[i])) {
             VLOG(6) << "SetInplaceOutputCorrectDistAttr Reshard inplace output"
                     << " to origin dist_attr "
                     << ReshardDebugInfo(*dist_tensor, dist_attr[i]);
             auto* func = phi::distributed::ChooseProperReshardFunction(
                 *dist_tensor, dist_attr[i]);
             func->Eval(dev_ctx, *dist_tensor, dist_attr[i], dist_tensor);
-          } else {
-            // just set correct SPMD dist_attrs
-            VLOG(6) << "SetInplaceOutputCorrectDistAttr input "
-                    << tensors[i].name() << " set its dist_attr from "
-                    << dist_tensor->dist_attr() << " to " << dist_attr[i];
-            dist_tensor->unsafe_set_dist_attr(dist_attr[i]);
+            continue;
           }
         }
-      } else {
-        VLOG(6) << "SetInplaceOutputCorrectDistAttr has"
-                << " uninitialized DistTensor input " << tensors[i].name()
-                << ", just set its dist_attr from " << dist_tensor->dist_attr()
-                << " to " << dist_attr[i];
-        dist_tensor->unsafe_set_dist_attr(dist_attr[i]);
       }
+      VLOG(6) << "SetInplaceOutputCorrectDistAttr for tensor "
+              << tensors[i].name() << ", just set its dist_attr from "
+              << dist_tensor->dist_attr() << " to " << dist_attr[i];
+      dist_tensor->unsafe_set_dist_attr(dist_attr[i]);
     }
   }
 }
diff --git a/paddle/phi/infermeta/spmd_rules/optimizer.cc b/paddle/phi/infermeta/spmd_rules/optimizer.cc
@@ -17,6 +17,7 @@ limitations under the License. */
 #include "glog/logging.h"
 
 #include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_utils.h"
 #include "paddle/phi/core/distributed/auto_parallel/utils.h"
 #include "paddle/phi/infermeta/spmd_rules/elementwise.h"
 #include "paddle/phi/infermeta/spmd_rules/utils.h"
@@ -86,19 +87,32 @@ SpmdInfo AdamInferSpmdDynamic(const DistMetaTensor& param,
       master_param.initialized()
           ? CopyTensorDistAttrForOutput(master_param.dist_attr())
           : TensorDistAttr();
-  TensorDistAttr skip_update_dist_attr =
-      skip_update.initialized()
-          ? CopyTensorDistAttrForOutput(skip_update.dist_attr())
-          : TensorDistAttr();
-
+  // If skip_update is on global_mesh, it should be reshard into
+  // local mesh. (currently occurs in static mode pipeline parellel)
+  auto skip_update_dist_attr = TensorDistAttr();
+  if (skip_update.initialized()) {
+    skip_update_dist_attr = skip_update.dist_attr();
+    PADDLE_ENFORCE_EQ(
+        skip_update_dist_attr.dims_mapping()[0],
+        -1,
+        errors::InvalidArgument(
+            "skip_update should be replicated, but got shard on mesh %d.",
+            skip_update_dist_attr.dims_mapping()[0]));
+    PADDLE_ENFORCE_EQ(
+        skip_update_dist_attr.partial_status().size(),
+        0,
+        errors::InvalidArgument("skip_update should be replicated, but got "
+                                "patial status not empty"));
+    if (skip_update_dist_attr.process_mesh().ndim() > 1 &&
+        phi::distributed::IsSubMesh(skip_update_dist_attr.process_mesh(),
+                                    param_dist_attr.process_mesh())) {
+      skip_update_dist_attr.set_process_mesh(param_dist_attr.process_mesh());
+    }
+  }
   // set the unchanged dims mapping
   lr_dist_attr.set_dims_mapping(learning_rate.dist_attr().dims_mapping());
   beta1_pow_dist_attr.set_dims_mapping(beta1_pow.dist_attr().dims_mapping());
   beta2_pow_dist_attr.set_dims_mapping(beta2_pow.dist_attr().dims_mapping());
-  if (skip_update.initialized()) {
-    skip_update_dist_attr.set_dims_mapping(
-        skip_update.dist_attr().dims_mapping());
-  }
 
   // set the changeable dims mapping
   auto param_spmd_dims_mapping = param_dist_attr_spmd.dims_mapping();
diff --git a/python/paddle/distributed/auto_parallel/static/engine.py b/python/paddle/distributed/auto_parallel/static/engine.py
@@ -678,7 +678,7 @@ def _parallel_pir(self, mode):
                             decr_ratio=self._strategy.amp.decr_ratio,
                             use_dynamic_loss_scaling=self._strategy.amp.use_dynamic_loss_scaling,
                             use_amp_guard=self._strategy.amp.use_fp16_guard,
-                            use_master_grad=self._strategy.amp.use_dynamic_loss_scaling,
+                            use_master_grad=self._strategy.amp.use_master_grad,
                             use_promote=self._strategy.amp.use_promote,
                         )
                         # bfloat16 needs no scaler
@@ -1144,9 +1144,7 @@ def _initialize(self, mode, init_parameters=True):
             # 4. lazy init adaption
             # 5. amp init adaption
             # 6. vpp init adaption
-            self.program_helper.init_pir(
-                self._pir_dist_main_progs[mode], self._place
-            )
+
             # self._init_lr(self._pir_dense_main_progs[mode])
             self.program_helper.init_pir(
                 self._pir_dist_main_progs[mode], self._place
diff --git a/test/auto_parallel/pir/semi_auto_parallel_dist_to_static_mlp_pir.py b/test/auto_parallel/pir/semi_auto_parallel_dist_to_static_mlp_pir.py
@@ -64,6 +64,7 @@ def __init__(self):
         self._amp = eval(os.getenv("amp", '0'))
         self._master_weight = eval(os.getenv("use_master_weight", '0'))
         self._master_grad = eval(os.getenv("use_master_grad", '0'))
+        self._use_promote = eval(os.getenv("use_promote", '0'))
         self._amp_dtype = os.getenv("amp_dtype", 'float16')
         self._amp_level = os.getenv("amp_level", 'O0')
         self.mesh = dist.ProcessMesh([0, 1], dim_names=["x"])
@@ -102,6 +103,7 @@ def run_dy2static(self, layer, opt, dist_loader):
             amp.level = self._amp_level
             amp.use_master_weight = self._master_weight
             amp.use_master_weight = self._master_grad
+            amp.use_promote = self._use_promote
 
         # static training
         dist_model = dist.to_static(
@@ -166,6 +168,7 @@ def run_dynamic(self, layer, opt, dist_loader, is_recompute=False):
                     level=self._amp_level,
                     dtype=self._amp_dtype,
                     enable=self._amp,
+                    use_promote=self._use_promote,
                 ):
                     out = layer(image)
                     loss = loss_fn(out, label)
@@ -186,14 +189,14 @@ def test_mp_demo_net(self):
 
         self.set_random_seed(self._seed)
         dy_layer = DemoNet(self.mesh)
-        dy_opt = paddle.optimizer.SGD(
+        dy_opt = paddle.optimizer.Adam(
             learning_rate=0.01, parameters=dy_layer.parameters()
         )
 
         paddle.base.set_flags({'FLAGS_enable_pir_api': 1})
         self.set_random_seed(self._seed)
         dy2static_layer = DemoNet(self.mesh)
-        dy2static_opt = paddle.optimizer.SGD(
+        dy2static_opt = paddle.optimizer.Adam(
             learning_rate=0.01, parameters=dy2static_layer.parameters()
         )
         dist_dataloader = dist.shard_dataloader(
@@ -214,14 +217,14 @@ def test_dp_demo_net(self):
 
         self.set_random_seed(self._seed)
         dy_layer = DPDemoNet(self.mesh)
-        dy_opt = paddle.optimizer.SGD(
+        dy_opt = paddle.optimizer.Adam(
             learning_rate=0.01, parameters=dy_layer.parameters()
         )
 
         paddle.base.set_flags({'FLAGS_enable_pir_api': 1})
         self.set_random_seed(self._seed)
         dy2static_layer = DPDemoNet(self.mesh)
-        dy2static_opt = paddle.optimizer.SGD(
+        dy2static_opt = paddle.optimizer.Adam(
             learning_rate=0.01, parameters=dy2static_layer.parameters()
         )
         dist_dataloader = dist.shard_dataloader(
@@ -257,14 +260,14 @@ def test_pp_demo_net(self):
 
         self.set_random_seed(self._seed)
         dy_layer = PPDemoNet(mesh1, mesh2)
-        dy_opt = paddle.optimizer.SGD(
+        dy_opt = paddle.optimizer.Adam(
             learning_rate=0.01, parameters=dy_layer.parameters()
         )
 
         paddle.base.set_flags({'FLAGS_enable_pir_api': 1})
         self.set_random_seed(self._seed)
         dy2static_layer = PPDemoNet(mesh1, mesh2)
-        dy2static_opt = paddle.optimizer.SGD(
+        dy2static_opt = paddle.optimizer.Adam(
             learning_rate=0.01, parameters=dy2static_layer.parameters()
         )
         dist_dataloader = dist.shard_dataloader(
diff --git a/test/auto_parallel/pir/test_semi_auto_parallel_dist_to_static_pir.py b/test/auto_parallel/pir/test_semi_auto_parallel_dist_to_static_pir.py
@@ -49,9 +49,10 @@ def test_mlp_amp(self):
                 "backend": ["gpu"],
                 "amp": ['1'],
                 "amp_dtype": ['float16'],
-                'amp_level': ['O1'],
+                'amp_level': ['O1', 'O2'],
                 'use_master_weight': ['1'],
                 'use_master_grad': ['1'],
+                'use_promote': ['1'],
             },
         )
         for envs in envs_list:
