From 4e1da46413f370fbce7a3f1acc6d4ac57f8d2be9 Mon Sep 17 00:00:00 2001 From: ykarnati Date: Tue, 4 Aug 2026 14:25:50 -0700 Subject: [PATCH] Pass process groups to remaining GTP layers Signed-off-by: ykarnati --- .../models/common/embeddings/language_model_embedding.py | 4 ++++ megatron/core/models/hybrid/hybrid_model.py | 2 ++ megatron/core/ssm/mamba_mixer.py | 2 ++ megatron/core/transformer/moe/shared_experts.py | 9 +++++++-- 4 files changed, 15 insertions(+), 2 deletions(-) diff --git a/megatron/core/models/common/embeddings/language_model_embedding.py b/megatron/core/models/common/embeddings/language_model_embedding.py index 7e49ec6c02d..d5b6e952b41 100644 --- a/megatron/core/models/common/embeddings/language_model_embedding.py +++ b/megatron/core/models/common/embeddings/language_model_embedding.py @@ -6,6 +6,7 @@ from torch import Tensor from megatron.core import tensor_parallel +from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.transformer.module import MegatronModule from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.utils import get_tensor_model_parallel_group_if_none, nvtx_decorator @@ -24,6 +25,7 @@ class LanguageModelEmbedding(MegatronModule): num_tokentypes (int): Set to 0 without binary head, and 2 with a binary head. Defaults to 0. scatter_to_sequence_parallel (bool): Set to False to disable scatter of embedding across sequence parallel region. Defaults to True. + pg_collection (ProcessGroupCollection, optional): Process groups used by the embedding. """ def __init__( @@ -35,6 +37,7 @@ def __init__( num_tokentypes: int = 0, scatter_to_sequence_parallel: bool = True, tp_group: Optional[torch.distributed.ProcessGroup] = None, + pg_collection: Optional[ProcessGroupCollection] = None, ): super().__init__(config=config) @@ -60,6 +63,7 @@ def __init__( reduce_scatter_embeddings=self.reduce_scatter_embeddings, config=self.config, tp_group=self.tp_group, + pg_collection=pg_collection, ) # Position embedding (serial). diff --git a/megatron/core/models/hybrid/hybrid_model.py b/megatron/core/models/hybrid/hybrid_model.py index f750c77e05b..f0358de57b9 100644 --- a/megatron/core/models/hybrid/hybrid_model.py +++ b/megatron/core/models/hybrid/hybrid_model.py @@ -239,6 +239,7 @@ def __init__( position_embedding_type=position_embedding_type, scatter_to_sequence_parallel=scatter_embedding_sequence_parallel, tp_group=self.pg_collection.tp, + pg_collection=self.pg_collection, ) # MLA (also used by DeepSeek Sparse Attention) uses its own decoupled RoPE, therefore we do @@ -322,6 +323,7 @@ def __init__( skip_weight_param_allocation=self.pre_process and self.share_embeddings_and_output_weights, tp_group=self.pg_collection.tp, + pg_collection=self.pg_collection, ) if self.pre_process or self.post_process or self.mtp_process: diff --git a/megatron/core/ssm/mamba_mixer.py b/megatron/core/ssm/mamba_mixer.py index f6ae07dd230..73e0561fdbf 100644 --- a/megatron/core/ssm/mamba_mixer.py +++ b/megatron/core/ssm/mamba_mixer.py @@ -291,6 +291,7 @@ def __init__( is_expert=False, tp_comm_buffer_name="fc1", tp_group=self.pg_collection.tp, + pg_collection=self.pg_collection, name=(name + f".in_proj") if name is not None else None, ) # in_proj packs [z, x, B, C, dt] into one ColumnParallelLinear. Each @@ -442,6 +443,7 @@ def __init__( is_expert=False, tp_comm_buffer_name="fc2", tp_group=self.pg_collection.tp, + pg_collection=self.pg_collection, name=(name + f".out_proj") if name is not None else None, ) diff --git a/megatron/core/transformer/moe/shared_experts.py b/megatron/core/transformer/moe/shared_experts.py index 027d0a780ff..038a162f899 100644 --- a/megatron/core/transformer/moe/shared_experts.py +++ b/megatron/core/transformer/moe/shared_experts.py @@ -125,8 +125,13 @@ def __init__( "please set '--disable-bias-linear' instead." config.ffn_hidden_size = config.moe_shared_expert_intermediate_size - # TODO(Hepteract): pass pg_collection to MLP after refactoring MLP - super().__init__(config=config, submodules=submodules, tp_group=pg_collection.tp, name=name) + super().__init__( + config=config, + submodules=submodules, + tp_group=pg_collection.tp, + name=name, + pg_collection=pg_collection, + ) self.use_shared_expert_gate = gate if self.use_shared_expert_gate: