Merge branch 'llama3' into 'main'

Use TP-CP group for fp8 amax reduction

See merge request ADLR/megatron-lm!1747

Author: ericharper

megatron/core/parallel_state.py

@@ -887,14 +887,14 @@ def get_amax_reduction_group(with_context_parallel=False):
     """Get the FP8 amax reduction group the caller rank belongs to."""
     if with_context_parallel:
         assert (
-            _TENSOR_AND_DATA_PARALLEL_GROUP_WITH_CP is not None
+            _TENSOR_AND_CONTEXT_PARALLEL_GROUP is not None
         ), 'FP8 amax reduction group is not initialized'
-        return _TENSOR_AND_DATA_PARALLEL_GROUP_WITH_CP
+        return _TENSOR_AND_CONTEXT_PARALLEL_GROUP
     else:
         assert (
-            _TENSOR_AND_DATA_PARALLEL_GROUP is not None
+            _TENSOR_MODEL_PARALLEL_GROUP is not None
         ), 'FP8 amax reduction group is not initialized'
-        return _TENSOR_AND_DATA_PARALLEL_GROUP
+        return _TENSOR_MODEL_PARALLEL_GROUP
 
 
 def get_tensor_and_data_parallel_group(with_context_parallel=False):

tests/unit_tests/test_parallel_state.py

@@ -218,7 +218,7 @@ def test_different_initialize_order_unconsistency(src_tp_pp, ep_size):
     assert dp_g != torch.distributed.get_process_group_ranks(ps.get_data_parallel_group(False))
     assert pp_g != torch.distributed.get_process_group_ranks(ps.get_pipeline_model_parallel_group())
     assert cp_g == torch.distributed.get_process_group_ranks(ps.get_context_parallel_group())
-    assert amax_g != torch.distributed.get_process_group_ranks(ps.get_amax_reduction_group(False))
+    assert amax_g == torch.distributed.get_process_group_ranks(ps.get_amax_reduction_group(False))
     assert mp_g != torch.distributed.get_process_group_ranks(ps.get_model_parallel_group())
 
     Utils.destroy_model_parallel()