Pass explicit process groups to hybrid logging

megatron/core/models/hybrid/hybrid_layer_allocation.py

@@ -333,6 +333,8 @@ def select_pipeline_segment(
     vp_stage: Optional[int],
     first_stage_layers: Optional[int] = None,
     last_stage_layers: Optional[int] = None,
+    tp_group: Optional[torch.distributed.ProcessGroup] = None,
+    dp_cp_group: Optional[torch.distributed.ProcessGroup] = None,
 ) -> Tuple[List[str], int]:
     """Select and validate the pipeline segment for the given PP rank and VP stage.
 
@@ -352,6 +354,8 @@ def select_pipeline_segment(
             uneven PP. Only valid when the pattern has no pipe separators.
         last_stage_layers: Number of layers on the last pipeline stage for
             uneven PP. Only valid when the pattern has no pipe separators.
+        tp_group: Optional tensor-parallel process group used for per-stage logging.
+        dp_cp_group: Optional data/context-parallel process group used for per-stage logging.
 
     Returns:
         Tuple of (layer_type_list, layer_offset) where layer_type_list is
@@ -445,6 +449,8 @@ def select_pipeline_segment(
             f"HybridModel: pp_rank={pp_rank}/{pp_size}, vp_stage={vp_stage}, "
             f"layers='{''.join(selected)}' ({len(selected)} layers), "
             f"layer_offset={offset} (auto-split)",
+            tp_group=tp_group,
+            dp_cp_group=dp_cp_group,
         )
         return selected, offset
 
@@ -479,6 +485,8 @@ def select_pipeline_segment(
         f"segment_index={segment_index}/{len(segments)}, "
         f"layers='{my_segment}' ({len(layer_type_list)} layers), "
         f"layer_offset={layer_offset}",
+        tp_group=tp_group,
+        dp_cp_group=dp_cp_group,
     )
 
     return layer_type_list, layer_offset

megatron/core/models/hybrid/hybrid_model.py

@@ -40,6 +40,18 @@
 logger = logging.getLogger(__name__)
 
 
+def _hybrid_logging_pg_kwargs(pg_collection: ProcessGroupCollection) -> dict:
+    tp_group = getattr(pg_collection, 'tp', None)
+    dp_cp_group = getattr(pg_collection, 'dp_cp', None)
+    if (tp_group is None) != (dp_cp_group is None):
+        raise ValueError(
+            "pg_collection.tp and pg_collection.dp_cp must both be set or both be unset."
+        )
+    if tp_group is None:
+        return {}
+    return {'tp_group': tp_group, 'dp_cp_group': dp_cp_group}
+
+
 class HybridModel(LanguageModule, GraphableMegatronModule):
     """Hybrid language model.
 
@@ -188,12 +200,15 @@ def __init__(
         self.mtp_pattern = parsed.mtp_pattern
         self.mtp_num_depths = parsed.mtp_num_depths
 
+        logging_pg_kwargs = _hybrid_logging_pg_kwargs(self.pg_collection)
+
         layer_type_list, layer_offset = select_pipeline_segment(
             parsed.main_pattern or '',
             self.pg_collection.pp,
             vp_stage,
             first_stage_layers=self.config.num_layers_in_first_pipeline_stage,
             last_stage_layers=self.config.num_layers_in_last_pipeline_stage,
+            **logging_pg_kwargs,
         )
 
         # Determine if MTP is needed (based on pattern parsing)

tests/unit_tests/models/test_hybrid_model.py

@@ -3,6 +3,7 @@
 import os
 from datetime import timedelta
 from itertools import accumulate
+from types import SimpleNamespace
 
 import pytest
 import torch
@@ -18,7 +19,7 @@
 from megatron.core.inference.utils import InferenceMode
 from megatron.core.models.common.embeddings.yarn_rotary_pos_embedding import YarnRotaryEmbedding
 from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec
-from megatron.core.models.hybrid.hybrid_model import HybridModel
+from megatron.core.models.hybrid.hybrid_model import HybridModel, _hybrid_logging_pg_kwargs
 from megatron.core.packed_seq_params import PackedSeqParams
 from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
 from megatron.core.transformer import TransformerConfig
@@ -28,6 +29,118 @@
 from tests.unit_tests.test_utilities import Utils
 
 
+def test_hybrid_logging_process_groups_are_paired():
+    tp_group = object()
+    dp_cp_group = object()
+
+    assert _hybrid_logging_pg_kwargs(SimpleNamespace()) == {}
+    assert _hybrid_logging_pg_kwargs(SimpleNamespace(tp=tp_group, dp_cp=dp_cp_group)) == {
+        'tp_group': tp_group,
+        'dp_cp_group': dp_cp_group,
+    }
+
+    with pytest.raises(ValueError, match="tp.*dp_cp"):
+        _hybrid_logging_pg_kwargs(SimpleNamespace(tp=tp_group))
+    with pytest.raises(ValueError, match="tp.*dp_cp"):
+        _hybrid_logging_pg_kwargs(SimpleNamespace(dp_cp=dp_cp_group))
+    with pytest.raises(ValueError, match="tp.*dp_cp"):
+        _hybrid_logging_pg_kwargs(SimpleNamespace(tp=tp_group, dp_cp=None))
+    with pytest.raises(ValueError, match="tp.*dp_cp"):
+        _hybrid_logging_pg_kwargs(SimpleNamespace(tp=None, dp_cp=dp_cp_group))
+
+
+@pytest.mark.skipif(
+    not is_torch_min_version("2.4.0"),
+    reason="torch.distributed.init_device_mesh requires torch >= 2.4.0",
+)
+@pytest.mark.parametrize("tp_size,cp_size,pp_size", [(2, 1, 4), (1, 1, 8), (8, 1, 1)])
+def test_hybrid_model_with_custom_process_groups(tmp_path, tp_size, cp_size, pp_size):
+    """Test HybridModel with custom process groups."""
+    Utils.initialize_model_parallel(
+        tensor_model_parallel_size=tp_size,
+        context_parallel_size=cp_size,
+        pipeline_model_parallel_size=pp_size,
+    )
+
+    try:
+        # Create device mesh for custom process groups
+        assert torch.distributed.get_world_size() == 8, "Test requires 8 GPUs"
+
+        # Initialize torch.distributed if not already initialized
+        if not torch.distributed.is_initialized():
+            torch.distributed.init_process_group(backend='nccl')
+
+        dp_size = 1
+
+        # Create HyperCommGrid with dimensions tp, cp, pp, dp.
+        grid = HyperCommGrid([tp_size, cp_size, pp_size, dp_size], ["tp", "cp", "pp", "dp"])
+
+        pp_group = grid.create_pg("pp")
+        cp_group = grid.create_pg("cp")
+        tp_group = grid.create_pg("tp")
+        dp_cp_group = grid.create_pg(["cp", "dp"])
+        embd_group_ranks = parallel_state.default_embedding_ranks(
+            torch.distributed.get_process_group_ranks(pp_group)
+        )
+        embd_group = torch.distributed.new_group(
+            ranks=embd_group_ranks, timeout=timedelta(minutes=30)
+        )
+
+        # Create model with custom process groups
+        from megatron.core.process_groups_config import ProcessGroupCollection
+
+        pg_collection = ProcessGroupCollection(
+            tp=tp_group, cp=cp_group, pp=pp_group, embd=embd_group, dp_cp=dp_cp_group
+        )
+
+        # Build pattern with '|' pipeline stage separators: 2 layers per PP stage
+        hybrid_layer_pattern = "|".join(["*-"] * pp_size)
+
+        # Configure model with appropriate sizes for parallelism
+        model_config = TransformerConfig(
+            num_layers=2 * pp_size,  # Scale layers with PP size
+            hidden_size=256 * tp_size,
+            num_attention_heads=4 * tp_size,  # Scale heads with TP size
+            use_cpu_initialization=True,
+            tensor_model_parallel_size=tp_size,
+            context_parallel_size=cp_size,
+            pipeline_model_parallel_size=pp_size,
+            pipeline_dtype=torch.bfloat16,
+        )
+
+        model = HybridModel(
+            config=model_config,
+            hybrid_stack_spec=hybrid_stack_spec,
+            vocab_size=128,
+            max_sequence_length=4,
+            hybrid_layer_pattern=hybrid_layer_pattern,
+            pg_collection=pg_collection,
+        )
+
+        # Basic forward test
+        micro_batch_size = 2
+        sequence_length = model.max_sequence_length
+
+        model.cuda()
+
+        data = list(range(sequence_length))
+        input_ids = torch.tensor(data, dtype=torch.int64).repeat((micro_batch_size, 1)).cuda()
+        position_ids = torch.tensor(data, dtype=torch.int64).repeat((micro_batch_size, 1)).cuda()
+        attention_mask = torch.ones(
+            (micro_batch_size, 1, sequence_length, sequence_length), dtype=bool
+        ).cuda()
+
+        logits = model.forward(
+            input_ids=input_ids, position_ids=position_ids, attention_mask=attention_mask
+        )
+
+        assert logits.shape[0] == micro_batch_size
+        assert logits.shape[1] == sequence_length
+        assert logits.shape[2] == divide(model.vocab_size, tp_size)
+    finally:
+        Utils.destroy_model_parallel()
+
+
 class TestHybridModel:
 
     def setup_method(self, method):
@@ -213,91 +326,6 @@ def test_layer_numbers(self):
         for expected, layer in enumerate(model.decoder.layers, start=1):
             assert expected == layer.layer_number, "layer numbers are incorrect"
 
-    @pytest.mark.skipif(
-        not is_torch_min_version("2.4.0"),
-        reason="torch.distributed.init_device_mesh requires torch >= 2.4.0",
-    )
-    @pytest.mark.parametrize("tp_size,cp_size,pp_size", [(2, 1, 4), (1, 1, 8), (8, 1, 1)])
-    def test_with_custom_process_groups(self, tmp_path, tp_size, cp_size, pp_size):
-        """Test HybridModel with custom process groups."""
-        Utils.initialize_model_parallel(
-            tensor_model_parallel_size=tp_size,
-            context_parallel_size=cp_size,
-            pipeline_model_parallel_size=pp_size,
-        )
-
-        # Create device mesh for custom process groups
-        assert torch.distributed.get_world_size() == 8, "Test requires 8 GPUs"
-
-        # Initialize torch.distributed if not already initialized
-        if not torch.distributed.is_initialized():
-            torch.distributed.init_process_group(backend='nccl')
-
-        # Create HyperCommGrid with dimensions tp, cp, pp (reversed from device mesh order)
-        grid = HyperCommGrid([tp_size, cp_size, pp_size], ["tp", "cp", "pp"])
-
-        pp_group = grid.create_pg("pp")
-        cp_group = grid.create_pg("cp")
-        tp_group = grid.create_pg("tp")
-        embd_group_ranks = parallel_state.default_embedding_ranks(
-            torch.distributed.get_process_group_ranks(pp_group)
-        )
-        embd_group = torch.distributed.new_group(
-            ranks=embd_group_ranks, timeout=timedelta(minutes=30)
-        )
-
-        # Create model with custom process groups
-        from megatron.core.process_groups_config import ProcessGroupCollection
-
-        pg_collection = ProcessGroupCollection(
-            tp=tp_group, cp=cp_group, pp=pp_group, embd=embd_group
-        )
-
-        # Build pattern with '|' pipeline stage separators: 3 layers per PP stage
-        hybrid_layer_pattern = "|".join(["M*-"] * pp_size)
-
-        # Configure model with appropriate sizes for parallelism
-        model_config = TransformerConfig(
-            num_layers=3 * pp_size,  # Scale layers with PP size
-            hidden_size=256 * tp_size,
-            num_attention_heads=4 * tp_size,  # Scale heads with TP size
-            use_cpu_initialization=True,
-            tensor_model_parallel_size=tp_size,
-            context_parallel_size=cp_size,
-            pipeline_model_parallel_size=pp_size,
-            pipeline_dtype=torch.bfloat16,
-        )
-
-        model = HybridModel(
-            config=model_config,
-            hybrid_stack_spec=hybrid_stack_spec,
-            vocab_size=128,
-            max_sequence_length=4,
-            hybrid_layer_pattern=hybrid_layer_pattern,
-            pg_collection=pg_collection,
-        )
-
-        # Basic forward test
-        micro_batch_size = 2
-        sequence_length = model.max_sequence_length
-
-        model.cuda()
-
-        data = list(range(sequence_length))
-        input_ids = torch.tensor(data, dtype=torch.int64).repeat((micro_batch_size, 1)).cuda()
-        position_ids = torch.tensor(data, dtype=torch.int64).repeat((micro_batch_size, 1)).cuda()
-        attention_mask = torch.ones(
-            (micro_batch_size, 1, sequence_length, sequence_length), dtype=bool
-        ).cuda()
-
-        logits = model.forward(
-            input_ids=input_ids, position_ids=position_ids, attention_mask=attention_mask
-        )
-
-        assert logits.shape[0] == micro_batch_size
-        assert logits.shape[1] == sequence_length
-        assert logits.shape[2] == divide(model.vocab_size, tp_size)
-
 
 class TestHybridQKLayernorm:
 

tests/unit_tests/ssm/test_hybrid_layer_allocation.py

@@ -466,6 +466,16 @@ def test_logging_is_called(self, mock_log):
         select_pipeline_segment("M*M*", pp_group=None, vp_stage=None)
         mock_log.assert_called_once()
 
+    @patch('megatron.core.models.hybrid.hybrid_layer_allocation.log_on_each_pipeline_stage')
+    def test_logging_receives_explicit_groups(self, mock_log):
+        tp_group = object()
+        dp_cp_group = object()
+        select_pipeline_segment(
+            "M*M*", pp_group=None, vp_stage=None, tp_group=tp_group, dp_cp_group=dp_cp_group
+        )
+        assert mock_log.call_args.kwargs["tp_group"] is tp_group
+        assert mock_log.call_args.kwargs["dp_cp_group"] is dp_cp_group
+
     @patch('megatron.core.models.hybrid.hybrid_layer_allocation.log_on_each_pipeline_stage')
     def test_mutual_exclusivity_pipes_with_first_stage(self, mock_log):
         """Pipe separators + first_stage_layers should raise ValueError."""