fix(parallel.py): fix norm and moe gate gradient reduce check (#420)

Author: huangting4201

internlm/solver/optimizer/hybrid_zero_optim.py

@@ -23,7 +23,7 @@
 )
 from internlm.core.parallel.comm.isp import ISPCommunicatorWrapper
 from internlm.core.parallel.comm.zero import ParamAsyncBcastHandler
-from internlm.model.modules.utils import is_gate_param, is_moe_param
+from internlm.model.modules.utils import is_moe_param
 from internlm.monitor import send_alert_message
 from internlm.solver.optimizer.store import (
     BucketStore,
@@ -44,7 +44,7 @@
 from internlm.utils.common import get_current_device
 from internlm.utils.logger import get_logger
 from internlm.utils.megatron_timers import megatron_timer as timer
-from internlm.utils.parallel import is_using_isp, is_using_sequence_parallel
+from internlm.utils.parallel import is_using_isp, should_reduce_replica_param
 from internlm.utils.timeout import llm_timeout
 
 from .base_optimizer import BaseOptimizer
@@ -393,11 +393,7 @@ def extra_layernorm_reduce_grad_hook(*args):  # pylint: disable=W0613
                     # the grad of layernorm should be all-reduce across the global process group
                     # here is the first stage all-reduce in tp/wp process group
                     # the second stage all-reduce will be processed in reduce_grad_hook
-                    if (
-                        is_using_sequence_parallel()
-                        and hasattr(param, IS_REPLICA_ZERO_PARALLEL)
-                        and getattr(param, IS_REPLICA_ZERO_PARALLEL) is True
-                    ) or (is_gate_param(param) and gpc.config.parallel.expert.no_tp):
+                    if should_reduce_replica_param(param):
                         accum_grad_obj.register_hook(extra_layernorm_reduce_grad_hook)
 
                     # we should not only register for parameters which have isp_reduce_scatter_name attr.

internlm/utils/parallel.py

@@ -13,6 +13,7 @@
     ParallelMode,
 )
 from internlm.core.context import global_context as gpc
+from internlm.model.modules.utils import is_gate_param
 from internlm.utils.utils import TensorParallelMode
 
 
@@ -85,6 +86,35 @@ def is_replica_expert_data_parallel_parameter(p):
     return hasattr(p, IS_REPLICA_EXPERT_DATA_PARALLEL) and getattr(p, IS_REPLICA_EXPERT_DATA_PARALLEL)
 
 
+def should_reduce_replica_param(p):
+    _reduce = False
+
+    if not is_replica_zero_parallel_parameter(p):
+        return _reduce
+
+    # for replica parameter
+    if gpc.config.parallel["tensor"].get("mode", TensorParallelMode.mtp.name) == TensorParallelMode.mtp.name:
+        _reduce = False
+    elif gpc.config.parallel["tensor"].get("mode", TensorParallelMode.mtp.name) in (
+        TensorParallelMode.msp.name,
+        TensorParallelMode.fsp.name,
+    ):
+        _reduce = gpc.is_using_parallel_mode(ParallelMode.TENSOR)
+    elif gpc.config.parallel["tensor"].get("mode", TensorParallelMode.mtp.name) == TensorParallelMode.isp.name:
+        _reduce = gpc.is_using_parallel_mode(ParallelMode.WEIGHT)
+
+    if not is_gate_param(p):
+        return _reduce
+
+    # for moe gate parameter
+    if gpc.config.parallel["tensor"].get("mode", TensorParallelMode.mtp.name) == TensorParallelMode.mtp.name:
+        _reduce = gpc.is_using_parallel_mode(ParallelMode.TENSOR) and getattr(
+            gpc.config.parallel.expert, "no_tp", False
+        )
+
+    return _reduce
+
+
 def sync_model_param(model):
     r"""Make sure data parameters are consistent during Data Parallel Mode.