Use a separate communicator for AMAX reduction (NVIDIA#1585)

* add amax reduction group for fp8 training

* use a separate communicator for amax reduction across DP-ranks

reflect suggestion

Author: erhoo82

apex/transformer/parallel_state.py

@@ -48,6 +48,8 @@
 _DECODER_RELATIVE_POSITION_EMBEDDING_GROUP = None
 # Data parallel group that the current rank belongs to.
 _DATA_PARALLEL_GROUP = None
+# Data parallel AMAX reduction group that the current rank belongs to.
+_DATA_PARALLEL_GROUP_FOR_AMAX_REDUCTION = None
 
 _VIRTUAL_PIPELINE_MODEL_PARALLEL_RANK = None
 _VIRTUAL_PIPELINE_MODEL_PARALLEL_WORLD_SIZE = None
@@ -84,6 +86,7 @@ def initialize_model_parallel(
     pipeline_model_parallel_size_: int = 1,
     virtual_pipeline_model_parallel_size_: Optional[int] = None,
     pipeline_model_parallel_split_rank_: Optional[int] = None,
+    use_fp8_: bool = False,
     *,
     default_backend: Optional[str] = None,
     p2p_backend: Optional[str] = None,
@@ -96,6 +99,7 @@ def initialize_model_parallel(
         pipeline_model_parallel_size: number of GPUs used to parallelize model pipeline.
         virtual_pipeline_model_parallel_size: number of virtual stages (interleaved pipeline).
         pipeline_model_parallel_split_rank: for models with both encoder and decoder, rank in pipeline with split point.
+        use_fp8_: FP8 training that needs AMAX reduction across data-parallel ranks.
     Keyword Arguments:
         default_backend: Backend of process groups except for pipeline parallel ones.
             If :obj:`None`, the backend specified in `torch.distributed.init_process_group` will be used.
@@ -186,6 +190,9 @@ def initialize_model_parallel(
     # Build the data-parallel groups.
     global _DATA_PARALLEL_GROUP
     assert _DATA_PARALLEL_GROUP is None, "data parallel group is already initialized"
+    if use_fp8_:
+        global _DATA_PARALLEL_GROUP_FOR_AMAX_REDUCTION
+        assert _DATA_PARALLEL_GROUP_FOR_AMAX_REDUCTION is None, "amax reduction group is already initialized"
     all_data_parallel_group_ranks = []
     for i in range(pipeline_model_parallel_size):
         start_rank = i * num_pipeline_model_parallel_groups
@@ -196,6 +203,10 @@ def initialize_model_parallel(
             group = torch.distributed.new_group(ranks, backend=default_backend)
             if rank in ranks:
                 _DATA_PARALLEL_GROUP = group
+            if use_fp8_:
+                group = torch.distributed.new_group(ranks, backend=default_backend)
+                if rank in ranks:
+                    _DATA_PARALLEL_GROUP_FOR_AMAX_REDUCTION = group
 
     # Build the model-parallel groups.
     global _MODEL_PARALLEL_GROUP
@@ -363,6 +374,13 @@ def get_data_parallel_group():
     return _DATA_PARALLEL_GROUP
 
 
+def get_data_parallel_amax_reduction_group():
+    """Get the amax reduction group the caller rank belongs to."""
+    assert _DATA_PARALLEL_GROUP_FOR_AMAX_REDUCTION is not None, \
+        "AMAX reduction group is not initialized"
+    return _DATA_PARALLEL_GROUP_FOR_AMAX_REDUCTION
+
+
 def get_embedding_group():
     """Get the embedding group the caller rank belongs to."""
     assert _EMBEDDING_GROUP is not None, "embedding group is not initialized"
@@ -655,6 +673,8 @@ def destroy_model_parallel():
     _PIPELINE_MODEL_PARALLEL_GROUP = None
     global _DATA_PARALLEL_GROUP
     _DATA_PARALLEL_GROUP = None
+    global _DATA_PARALLEL_GROUP_FOR_AMAX_REDUCTION
+    _DATA_PARALLEL_GROUP_FOR_AMAX_REDUCTION = None
     global _EMBEDDING_GROUP
     _EMBEDDING_GROUP = None
     global _POSITION_EMBEDDING_GROUP