From c74185c2be0d2e64ec7260057d48cb196a4d2377 Mon Sep 17 00:00:00 2001 From: Yashaswi Karnati Date: Wed, 13 May 2026 02:10:14 +0000 Subject: [PATCH 1/5] Pass hybrid logging process groups --- .../core/models/hybrid/hybrid_layer_allocation.py | 8 ++++++++ megatron/core/models/hybrid/hybrid_model.py | 11 +++++++++++ tests/unit_tests/ssm/test_hybrid_layer_allocation.py | 10 ++++++++++ 3 files changed, 29 insertions(+) diff --git a/megatron/core/models/hybrid/hybrid_layer_allocation.py b/megatron/core/models/hybrid/hybrid_layer_allocation.py index f1ba94ef7fa..67103fe67f1 100644 --- a/megatron/core/models/hybrid/hybrid_layer_allocation.py +++ b/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 diff --git a/megatron/core/models/hybrid/hybrid_model.py b/megatron/core/models/hybrid/hybrid_model.py index 4b5858ef9da..ff5c92912e1 100644 --- a/megatron/core/models/hybrid/hybrid_model.py +++ b/megatron/core/models/hybrid/hybrid_model.py @@ -186,12 +186,23 @@ def __init__( self.mtp_pattern = parsed.mtp_pattern self.mtp_num_depths = parsed.mtp_num_depths + logging_pg_kwargs = {} + if ( + getattr(self.pg_collection, 'tp', None) is not None + and getattr(self.pg_collection, 'dp_cp', None) is not None + ): + logging_pg_kwargs = { + 'tp_group': self.pg_collection.tp, + 'dp_cp_group': self.pg_collection.dp_cp, + } + 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) diff --git a/tests/unit_tests/ssm/test_hybrid_layer_allocation.py b/tests/unit_tests/ssm/test_hybrid_layer_allocation.py index fe0d7c2dc1e..faa553216da 100644 --- a/tests/unit_tests/ssm/test_hybrid_layer_allocation.py +++ b/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.""" From fbcb2b357f7adb41dc591913c970fbea4f3c336f Mon Sep 17 00:00:00 2001 From: Yashaswi Karnati Date: Wed, 13 May 2026 02:28:59 +0000 Subject: [PATCH 2/5] Guard hybrid logging process groups --- megatron/core/models/hybrid/hybrid_model.py | 22 +++++++++++-------- tests/unit_tests/models/test_hybrid_model.py | 23 +++++++++++++++++++- 2 files changed, 35 insertions(+), 10 deletions(-) diff --git a/megatron/core/models/hybrid/hybrid_model.py b/megatron/core/models/hybrid/hybrid_model.py index ff5c92912e1..22d96aad3b1 100644 --- a/megatron/core/models/hybrid/hybrid_model.py +++ b/megatron/core/models/hybrid/hybrid_model.py @@ -38,6 +38,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. @@ -186,15 +198,7 @@ def __init__( self.mtp_pattern = parsed.mtp_pattern self.mtp_num_depths = parsed.mtp_num_depths - logging_pg_kwargs = {} - if ( - getattr(self.pg_collection, 'tp', None) is not None - and getattr(self.pg_collection, 'dp_cp', None) is not None - ): - logging_pg_kwargs = { - 'tp_group': self.pg_collection.tp, - 'dp_cp_group': self.pg_collection.dp_cp, - } + logging_pg_kwargs = _hybrid_logging_pg_kwargs(self.pg_collection) layer_type_list, layer_offset = select_pipeline_segment( parsed.main_pattern or '', diff --git a/tests/unit_tests/models/test_hybrid_model.py b/tests/unit_tests/models/test_hybrid_model.py index 98a53da0314..ddec24c1f5a 100644 --- a/tests/unit_tests/models/test_hybrid_model.py +++ b/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 @@ -17,7 +18,7 @@ from megatron.core.inference.sampling_params import SamplingParams 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 @@ -27,6 +28,26 @@ 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)) + + class TestHybridModel: def setup_method(self, method): From 7804d000f67be679ef6adff55d3659f7f696b563 Mon Sep 17 00:00:00 2001 From: ykarnati Date: Mon, 1 Jun 2026 15:14:54 -0700 Subject: [PATCH 3/5] Fix hybrid custom process group test --- tests/unit_tests/models/test_hybrid_model.py | 175 ++++++++++--------- 1 file changed, 89 insertions(+), 86 deletions(-) diff --git a/tests/unit_tests/models/test_hybrid_model.py b/tests/unit_tests/models/test_hybrid_model.py index ddec24c1f5a..ca9ad134a55 100644 --- a/tests/unit_tests/models/test_hybrid_model.py +++ b/tests/unit_tests/models/test_hybrid_model.py @@ -48,6 +48,95 @@ def test_hybrid_logging_process_groups_are_paired(): _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') + + # 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, dp_cp=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): @@ -228,92 +317,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: def setup_method(self, method): From e2bc62d18ce21491a02cc7296c4475ed3a59e3b1 Mon Sep 17 00:00:00 2001 From: ykarnati Date: Mon, 1 Jun 2026 16:50:47 -0700 Subject: [PATCH 4/5] Use explicit dp_cp group in hybrid test --- tests/unit_tests/models/test_hybrid_model.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/tests/unit_tests/models/test_hybrid_model.py b/tests/unit_tests/models/test_hybrid_model.py index ca9ad134a55..26a7ccff179 100644 --- a/tests/unit_tests/models/test_hybrid_model.py +++ b/tests/unit_tests/models/test_hybrid_model.py @@ -69,12 +69,17 @@ def test_hybrid_model_with_custom_process_groups(tmp_path, tp_size, cp_size, pp_ 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"]) + 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) ) @@ -86,7 +91,7 @@ def test_hybrid_model_with_custom_process_groups(tmp_path, tp_size, cp_size, pp_ 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=cp_group + 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 From 76ad15a2a3718895c72e0be0d40a6db014391a1a Mon Sep 17 00:00:00 2001 From: ykarnati Date: Mon, 1 Jun 2026 17:38:24 -0700 Subject: [PATCH 5/5] Format hybrid process group test --- tests/unit_tests/models/test_hybrid_model.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tests/unit_tests/models/test_hybrid_model.py b/tests/unit_tests/models/test_hybrid_model.py index 078dbb4ebd5..ffc9fe41e99 100644 --- a/tests/unit_tests/models/test_hybrid_model.py +++ b/tests/unit_tests/models/test_hybrid_model.py @@ -73,9 +73,7 @@ def test_hybrid_model_with_custom_process_groups(tmp_path, tp_size, cp_size, pp_ 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"] - ) + 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") @@ -328,6 +326,7 @@ def test_layer_numbers(self): for expected, layer in enumerate(model.decoder.layers, start=1): assert expected == layer.layer_number, "layer numbers are incorrect" + class TestHybridQKLayernorm: def setup_method(self, method):