Fix comments for PR #59644 (#59885)

wentaoyu · web-flow · commit 843d10184d58 · 2023-12-12T19:20:58.000+08:00
* update

* update
diff --git a/python/paddle/distributed/fleet/meta_optimizers/dygraph_optimizer/hybrid_parallel_optimizer.py b/python/paddle/distributed/fleet/meta_optimizers/dygraph_optimizer/hybrid_parallel_optimizer.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import distutils.util
 import os
 
 import paddle
@@ -46,6 +47,9 @@ def __init__(self, clip, hcg):
         self._clip = clip
         self._hcg = hcg
         self.not_sharding_stage1 = True
+        self._force_align_vpp_grad_sum_order = distutils.util.strtobool(
+            os.getenv('FLAGS_force_align_vpp_grad_sum_order', '0')
+        )
 
     def _global_norm(self, global_norm_var_dist, global_norm_var_not_dist):
         # sharding first
@@ -99,6 +103,10 @@ def _global_norm(self, global_norm_var_dist, global_norm_var_not_dist):
 
     @no_grad()
     def _dygraph_clip(self, params_grads):
+        if self._force_align_vpp_grad_sum_order:
+            chunk_num = self._get_vpp_chunk_num(params_grads)
+            if chunk_num > 0:
+                return self._vpp_dygraph_clip(params_grads, chunk_num)
         sum_square_dist_fp16 = []
         sum_square_dist_bf16 = []
         sum_square_dist_fp32 = []
diff --git a/python/paddle/distributed/fleet/meta_parallel/pipeline_parallel.py b/python/paddle/distributed/fleet/meta_parallel/pipeline_parallel.py
@@ -921,8 +921,6 @@ def __init__(self, layers, hcg, strategy):
         self._virtual_pp_rank = 0
         self._reset_counter()
 
-        self._check_sanity()
-
     def _check_sanity(self):
         assert (
             framework.in_dynamic_mode()
diff --git a/python/paddle/distributed/fleet/utils/tensor_fusion_helper.py b/python/paddle/distributed/fleet/utils/tensor_fusion_helper.py
@@ -461,6 +461,15 @@ def scale_grads(self):
 
         self._reset_params_checked_in()
 
+    @imperative_base.no_grad
+    def scale_and_split_grads(self):
+        assert self._task is not None, "Task is not initialized. "
+        self._task.wait()
+        scale_factor = 1.0 / self._comm_group.nranks
+        self.grad_storage.scale_(scale_factor)
+
+        self._reset_params_checked_in()
+
 
 def obtain_storage(
     parameters,
