Pass process groups to remaining GTP-backed layers

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).

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:

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,
         )
 

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: