Require an explicit process-group collection in LanguageModule

examples/academic_paper_scripts/detxoify_lm/generate_samples_gpt.py

@@ -13,6 +13,7 @@
 from megatron.training import get_tokenizer
 from megatron.training import print_rank_0
 from megatron.training.checkpointing import load_checkpoint
+from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core import mpu
 from megatron.training.arguments import parse_and_validate_args
 from megatron.training.initialize import initialize_megatron
@@ -74,7 +75,9 @@ def model_provider(pre_process=True, post_process=True) -> GPTModel:
         share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights,
         position_embedding_type=args.position_embedding_type,
         rotary_percent=args.rotary_percent
-    )
+    ,
+                pg_collection=ProcessGroupCollection.use_mpu_process_groups(),
+            )
 
     return model
 

examples/bert/pretrain_bert.py

@@ -10,6 +10,7 @@
 from megatron.training import get_args
 from megatron.training import print_rank_0
 from megatron.training import get_timers
+from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core import tensor_parallel
 from megatron.core.enums import ModelType
 from megatron.core.models.bert.bert_model import BertModel
@@ -55,7 +56,9 @@ def model_provider(pre_process=True, post_process=True, vp_stage=None, config=No
         parallel_output=True,
         pre_process=pre_process,
         post_process=post_process,
-        vp_stage=vp_stage)
+        vp_stage=vp_stage,
+                pg_collection=ProcessGroupCollection.use_mpu_process_groups(),
+            )
 
     return model
 

examples/export/trtllm_export/distributed_export/gpt_distributed_gpu_export.py

@@ -1,5 +1,6 @@
 import os
 import torch
+from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core import parallel_state
 from megatron.core import dist_checkpointing
 from megatron.core.export.model_type import ModelType
@@ -42,7 +43,9 @@ def model_provider():
         transformer_layer_spec=get_gpt_layer_local_spec(), 
         vocab_size=_VOCAB_SIZE, 
         max_sequence_length=_SEQUENCE_LENGTH,
-    )
+
+                    pg_collection=ProcessGroupCollection.use_mpu_process_groups(),
+                )
 
     return gpt_model
 

examples/export/trtllm_export/single_device_export/gpt_single_device_cpu_export.py

@@ -1,5 +1,6 @@
 import os
 import torch
+from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core import parallel_state
 from megatron.core import dist_checkpointing
 from megatron.core.export.model_type import ModelType
@@ -43,7 +44,9 @@ def model_provider():
         transformer_layer_spec=get_gpt_layer_local_spec(), 
         vocab_size=100, 
         max_sequence_length=_SEQUENCE_LENGTH,
-    )
+
+                    pg_collection=ProcessGroupCollection.use_mpu_process_groups(),
+                )
 
     return gpt_model
 

examples/run_simple_mcore_train_loop.py

@@ -7,6 +7,7 @@
 from functools import partial
 from pathlib import Path
 from typing import Any, Callable, Dict, Tuple, Iterator
+from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core import parallel_state
 from megatron.core import dist_checkpointing
 from megatron.core.pipeline_parallel.schedules import get_forward_backward_func
@@ -74,7 +75,9 @@ def model_provider() -> GPTModel:
         transformer_layer_spec=get_gpt_layer_local_spec(),
         vocab_size=100,
         max_sequence_length=_SEQUENCE_LENGTH,
-    )
+
+                              pg_collection=ProcessGroupCollection.use_mpu_process_groups(),
+                          )
 
     return gpt_model
 

examples/t5/pretrain_t5.py

@@ -9,6 +9,7 @@
 import torch
 
 import megatron
+from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core import mpu, tensor_parallel
 from megatron.core.tokenizers.utils.build_tokenizer import build_tokenizer
 from megatron.core.datasets.blended_megatron_dataset_builder import BlendedMegatronDatasetBuilder
@@ -132,7 +133,9 @@ def model_provider(
         relative_attention_max_distance=args.relative_attention_max_distance,
         add_encoder=add_encoder,
         add_decoder=add_decoder,
-    )
+
+                pg_collection=ProcessGroupCollection.use_mpu_process_groups(),
+            )
 
     return model
 

megatron/core/models/T5/t5_model.py

@@ -158,7 +158,7 @@ def __init__(
         pg_collection: ProcessGroupCollection = None,
     ):
 
-        super(T5Model, self).__init__(config=config)
+        super(T5Model, self).__init__(config=config, pg_collection=pg_collection)
 
         self.config: TransformerConfig = config
         self.encoder_config: TransformerConfig = encoder_config

megatron/core/models/common/language_module/language_module.py

@@ -47,8 +47,12 @@ def __init__(
     ) -> None:
         super().__init__(config=config)
         self._set_attention_backend()
-        if pg_collection is None:
-            pg_collection = ProcessGroupCollection.use_mpu_process_groups()
+        assert pg_collection is not None, (
+            "LanguageModule requires an explicit pg_collection. The global parallel grid is not a "
+            "safe default: a model built on independent grids (vision encoder + LLM, GTP, MIMO) "
+            "would silently get the wrong one. "
+            "See docs/developer/parallel-state-deprecation.md"
+        )
         self.pg_collection = pg_collection
         self.cp_group = pg_collection.cp
         self.tp_group = get_tensor_model_parallel_group_if_none(pg_collection.tp)

megatron/elastification/pretrain_hybrid_flex.py

@@ -7,6 +7,7 @@
 
 import torch
 
+from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core import mpu, parallel_state
 from megatron.core.datasets.blended_megatron_dataset_builder import BlendedMegatronDatasetBuilder
 from megatron.core.datasets.gpt_dataset import GPTDataset, GPTDatasetConfig, MockGPTDataset
@@ -140,7 +141,9 @@ def model_provider(pre_process=True, post_process=True, vp_stage: Optional[int]
         rotary_percent=args.rotary_percent,
         rotary_base=args.rotary_base,
         vp_stage=vp_stage
-    )
+    ,
+                pg_collection=ProcessGroupCollection.use_mpu_process_groups(),
+            )
     from megatron.elastification.flextron_utils import (
         inject_flextron_forward_logic,
         setup_flextron_model,

tests/unit_tests/a2a_overlap/test_cuda_graphed_schedule_chunk_1f1b.py

@@ -6,6 +6,7 @@
 import pytest
 import torch
 
+from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.enums import ModelType
 from megatron.core.models.gpt.gpt_layer_specs import (
     get_gpt_decoder_block_spec,
@@ -110,7 +111,9 @@ def model_provider(
             position_embedding_type=args.position_embedding_type,
             rotary_percent=args.rotary_percent,
             mtp_block_spec=mtp_block_spec,
-        )
+
+                   pg_collection=ProcessGroupCollection.use_mpu_process_groups(),
+               )
 
     def create_test_args(
         self, cuda_graph_impl, cuda_graph_modules, cuda_graph_warmup_steps, ep_size, **kwargs

tests/unit_tests/a2a_overlap/test_schedule_chunk_1f1b.py

@@ -4,6 +4,7 @@
 import pytest
 import torch
 
+from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.models.common.model_chunk_schedule_plan import TransformerModelChunkSchedulePlan
 from megatron.core.models.gpt.gpt_layer_specs import (
     get_gpt_decoder_block_spec,
@@ -62,7 +63,9 @@ def build_model(config, use_padding_mask=False):
         pre_process=True,
         post_process=True,
         max_sequence_length=max_seq_len,
-    )
+
+                    pg_collection=ProcessGroupCollection.use_mpu_process_groups(),
+                )
     f_schedule_plan = gpt_model.build_schedule_plan(**data)
     return gpt_model, f_schedule_plan, data
 

tests/unit_tests/a2a_overlap/test_schedule_layer_1f1b.py

@@ -5,6 +5,7 @@
 import torch
 import torch.nn.functional as F
 
+from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.fp8_utils import get_fp8_context
 from megatron.core.models.common.model_chunk_schedule_plan import TransformerLayerSchedulePlan
 from megatron.core.models.gpt.gpt_layer_specs import (
@@ -304,7 +305,9 @@ def test_transformer_layer_overlap_dense(self):
                 pre_process=True,
                 post_process=True,
                 max_sequence_length=300,
-            )
+
+                            pg_collection=ProcessGroupCollection.use_mpu_process_groups(),
+                        )
 
             params = reset_model(gpt_model)
             input_tensors = [build_data() for _ in range(microbatches)]
@@ -347,7 +350,9 @@ def test_transformer_layer_overlap_shared_expert(self):
                 pre_process=True,
                 post_process=True,
                 max_sequence_length=300,
-            )
+
+                            pg_collection=ProcessGroupCollection.use_mpu_process_groups(),
+                        )
 
             params = reset_model(gpt_model)
             input_tensors = [build_data() for _ in range(microbatches)]
@@ -366,7 +371,9 @@ def test_transformer_layer_overlap_shared_expert(self):
                 pre_process=True,
                 post_process=True,
                 max_sequence_length=300,
-            )
+
+                            pg_collection=ProcessGroupCollection.use_mpu_process_groups(),
+                        )
             reset_model(gpt_model, params)
             capture_a2a_overlap = run_transformer_layer_a2a_overlap_with_capture(
                 gpt_model, input_tensors, microbatches
@@ -399,7 +406,9 @@ def test_transformer_layer_overlap_early_attn_memory_release(self):
                 pre_process=True,
                 post_process=True,
                 max_sequence_length=300,
-            )
+
+                            pg_collection=ProcessGroupCollection.use_mpu_process_groups(),
+                        )
 
             params = reset_model(gpt_model)
             input_tensors = [build_data() for _ in range(microbatches)]
@@ -418,7 +427,9 @@ def test_transformer_layer_overlap_early_attn_memory_release(self):
                 pre_process=True,
                 post_process=True,
                 max_sequence_length=300,
-            )
+
+                            pg_collection=ProcessGroupCollection.use_mpu_process_groups(),
+                        )
             reset_model(gpt_model, params)
             capture_a2a_overlap = run_transformer_layer_a2a_overlap_with_capture(
                 gpt_model, input_tensors, microbatches
@@ -452,7 +463,9 @@ def test_transformer_layer_overlap(self, dispatcher_type, flex_backend, fp8_flag
                 pre_process=True,
                 post_process=True,
                 max_sequence_length=300,
-            )
+
+                            pg_collection=ProcessGroupCollection.use_mpu_process_groups(),
+                        )
 
             params = reset_model(gpt_model)
             input_tensors = [build_data() for _ in range(microbatches)]
@@ -512,7 +525,9 @@ def test_transformer_layer_overlap_zero_copy(self):
                     pre_process=True,
                     post_process=True,
                     max_sequence_length=300,
-                )
+
+                                pg_collection=ProcessGroupCollection.use_mpu_process_groups(),
+                            )
                 params = reset_model(gpt_model)
                 input_tensors = [build_data() for _ in range(microbatches)]
 
@@ -584,7 +599,9 @@ def test_mtp_layer_overlap(self, dispatcher_type, flex_backend, fp8_flag):
                 pre_process=True,
                 post_process=True,
                 max_sequence_length=300,
-            )
+
+                            pg_collection=ProcessGroupCollection.use_mpu_process_groups(),
+                        )
             gpt_model.decoder.final_layernorm = None
             gpt_model.cuda()
             params = reset_model(gpt_model)

tests/unit_tests/a2a_overlap/utils.py

@@ -5,6 +5,7 @@
 
 import torch
 
+from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core import config
 from megatron.core.transformer.transformer_config import MLATransformerConfig
 from megatron.core.utils import is_te_min_version
@@ -303,7 +304,9 @@ def build_gpt_model(config, vocab_size=512, max_seq_len=300):
         pre_process=True,
         post_process=True,
         max_sequence_length=max_seq_len,
-    )
+
+                pg_collection=ProcessGroupCollection.use_mpu_process_groups(),
+            )
     model.cuda()
     return model
 

tests/unit_tests/determinism/correctness/test_gpt_model.py

@@ -16,6 +16,7 @@
 import pytest
 import torch
 
+from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec
 from megatron.core.models.gpt.gpt_model import GPTModel
 from megatron.core.transformer.transformer_config import TransformerConfig
@@ -41,7 +42,9 @@ def build_gpt(overrides, pre_process=True, post_process=True, vp_stage=None, **_
         post_process=post_process,
         vp_stage=vp_stage,
         position_embedding_type="rope",
-    ).cuda()
+
+               pg_collection=ProcessGroupCollection.use_mpu_process_groups(),
+           ).cuda()
 
 
 def make_gpt_inputs():

tests/unit_tests/determinism/correctness/test_hybrid_model.py

@@ -10,6 +10,7 @@
 import pytest
 import torch
 
+from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec
 from megatron.core.models.hybrid.hybrid_model import HybridModel
 from megatron.core.transformer.transformer_config import TransformerConfig
@@ -108,7 +109,9 @@ def build(overrides, pre_process=True, post_process=True, vp_stage=None, **_):
                 pre_process=pre_process,
                 post_process=post_process,
                 vp_stage=vp_stage,
-            ).cuda()
+
+                       pg_collection=ProcessGroupCollection.use_mpu_process_groups(),
+                   ).cuda()
 
         runner = BitExactRunner(
             build_model=build,

tests/unit_tests/dist_checkpointing/models/test_bert_model.py

@@ -5,6 +5,7 @@
 import pytest
 import torch
 
+from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core import parallel_state as ps
 from megatron.core.models.bert.bert_layer_specs import (
     bert_layer_local_spec,
@@ -49,7 +50,9 @@ def initialize_bert_model(
         pre_process=pre_process,
         post_process=post_process,
         num_tokentypes=0,
-    )
+
+                pg_collection=ProcessGroupCollection.use_mpu_process_groups(),
+            )
 
     with torch.no_grad():
         for p in model.parameters():