Merge branch 'hongbinl/support_pp_1' into 'main'

feat(MoE): Support Expert Parallel A2A Overlapping - (02) Support EP A2A overlap at PP=1

See merge request ADLR/megatron-lm!3470

Author: ko3n1g

megatron/core/model_parallel_config.py

@@ -237,6 +237,14 @@ class ModelParallelConfig:
        Set the bootstrapping backend out of 'nccl', 'mpi', and 'gloo'
     """
 
+    overlap_moe_expert_parallel_comm: bool = False
+    """Overlap EP A2A communications with independent computations of different micro-batches
+    in 1f1b phase of pipelining or non-pipelining schedule.
+    """
+
+    delay_wgrad_compute: bool = False
+    """Delay the weight gradient computation to improve batch-level communication overlapping"""
+
     ###################
     # Pipeline Parallel
     ###################
@@ -307,9 +315,6 @@ class ModelParallelConfig:
        rank 1 |   0 1 2 0 1 2 3 4 3 4
     """
 
-    delay_wgrad_compute: bool = False
-    """If true, delay the wgrad compute for better overlapping in combined 1F1B."""
-
     ###################
     # CPU Offloading
     ###################