Support sequence_parallel with main_grad

ppfleetx/models/language_model/gpt/dygraph/hybrid_model.py

@@ -771,7 +771,7 @@ def __init__(self,
 
         hcg = env.get_hcg()
         mp_size = hcg.get_model_parallel_world_size()
-        if use_flash_attn or mp_size <= 1:
+        if mp_size <= 1:
             sequence_parallel = False
             logging.warning(
                 "If mp_size <= 1, sequence_parallel strategy will be turned off in GPTModelHybrid model."

ppfleetx/models/language_model/gpt/dygraph/sequence_parallel_utils.py

@@ -14,6 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import numpy as np
 import paddle
 from paddle import framework
 from paddle import distributed as dist
@@ -169,18 +170,19 @@ def __impl__(grad):
     return __impl__
 
 
-def create_non_fused_allreduce_gradient_hook(accumulation_steps):
+def create_non_fused_allreduce_gradient_hook(param, accumulation_steps):
     hcg = env.get_hcg()
     pg = hcg.get_model_parallel_group().process_group
-
     step = [0]
 
-    def __impl__(grad):
+    @paddle.autograd.no_grad()
+    def __impl__():
         step[0] += 1
-        if step[0] == accumulation_steps:
-            step[0] = 0
-            pg.allreduce(grad).wait()
-        return grad
+        if (step[0] % accumulation_steps) == 0:
+            if hasattr(param, "main_grad"):
+                pg.allreduce(param.main_grad).wait()
+            else:
+                pg.allreduce(param.grad).wait()
 
     return __impl__
 
@@ -202,11 +204,12 @@ def register_sequence_parallel_allreduce_hooks(
     if fuse_sequence_parallel_allreduce:
         hook = create_fused_allreduce_gradient_hook(params, accumulation_steps)
         for p in params:
-            p.register_hook(hook)
+            p._register_backward_hook(hook)
     else:
         for p in params:
-            p.register_hook(
-                create_non_fused_allreduce_gradient_hook(accumulation_steps))
+            hook = create_non_fused_allreduce_gradient_hook(p,
+                                                            accumulation_steps)
+            p._register_backward_hook(hook)
 
 
 def is_fused_matmul_bias_supported():

ppfleetx/utils/config.py

@@ -94,6 +94,12 @@ def process_dist_config(configs):
     if 'fuse_sequence_parallel_allreduce' not in config:
         config['fuse_sequence_parallel_allreduce'] = False
 
+    if 'use_main_grad' in config and config['use_main_grad'] is True:
+        logger.warning(
+            "If use_main_grad is True, fuse_sequence_parallel_allreduce will be forced to False"
+        )
+        config['fuse_sequence_parallel_allreduce'] = False
+
 
 def process_global_configs(config):
     """