Check parameter hashes across the model's own data-parallel groups

megatron/core/utils.py

@@ -978,7 +978,11 @@ def log_on_each_pipeline_stage(
 
 
 def check_param_hashes_across_dp_replicas(
-    model: List[torch.nn.Module], cross_check: bool = False
+    model: List[torch.nn.Module],
+    cross_check: bool = False,
+    *,
+    dp_group: Optional[torch.distributed.ProcessGroup] = None,
+    expt_dp_group: Optional[torch.distributed.ProcessGroup] = None,
 ) -> bool:
     """Computes hashes of all parameters in model, all-gathers hashes across DP replicas,
     and then checks for equality between the locally-computed hashes and those of other ranks.
@@ -987,15 +991,34 @@ def check_param_hashes_across_dp_replicas(
     tensors from GPU to CPU first; as a result, this function is not intended to be called
     very frequently in the main training loop.
 
+    The DP replicas are given by the model's data-parallel groups. If neither group is passed,
+    the data-parallel groups of the global parallel state are used.
+
     Args:
         model (List[torch.nn.Module]): List of model chunks whose parameter hashes need to
             be checked.
         cross_check (bool): If true, will check whether hashes match across all DP replicas.
+        dp_group (torch.distributed.ProcessGroup, optional): The model's data-parallel
+            replicate group (``pg_collection.dp``), across which the hashes of non-expert
+            parameters are compared.
+        expt_dp_group (torch.distributed.ProcessGroup, optional): The model's expert
+            data-parallel replicate group (``pg_collection.expt_dp``), across which the hashes
+            of expert parameters (``param.allreduce`` is False) are compared. May be None for a
+            model without expert parameters.
 
     Returns:
         True if all param hashes match with corresponding hash on DP replica 0 or
         across all replicas if cross_check is enabled, False otherwise.
+
+    Raises:
+        ValueError: If a group is passed, but this rank holds non-expert parameters and
+            ``dp_group`` is None, or expert parameters and ``expt_dp_group`` is None.
     """
+    if dp_group is None and expt_dp_group is None:
+        # Migration fallback for callers that do not pass the model's groups: the replicate
+        # data-parallel groups of the global parallel state.
+        dp_group = parallel_state.get_data_parallel_group(with_gtp_remat=False)
+        expt_dp_group = parallel_state.get_expert_data_parallel_group(with_gtp_remat=False)
 
     # Compute per-parameter hashes on this rank.
     # Keep track of expert and non-expert parameters separately since they need to be
@@ -1017,16 +1040,26 @@ def check_param_hashes_across_dp_replicas(
                 expert_params.append((model_chunk_id, param_name, param))
                 local_expert_param_hashes.append(param_hash)
 
+    # Before the first collective, make sure there is a group for each kind of parameter that
+    # this rank holds.
+    if non_expert_params and dp_group is None:
+        raise ValueError(
+            "check_param_hashes_across_dp_replicas: dp_group is None, but the model has "
+            "non-expert parameters; pass the model's pg_collection.dp."
+        )
+    if expert_params and expt_dp_group is None:
+        raise ValueError(
+            "check_param_hashes_across_dp_replicas: expt_dp_group is None, but the model has "
+            "expert parameters; pass the model's pg_collection.expt_dp."
+        )
+
     # Use data-modulo-expert parallel group to all-gather expert param hashes, regular
     # data-parallel group for non-expert param hashes.
     all_param_hashes_match = True
     for params, local_param_hashes, all_gather_group in zip(
         [non_expert_params, expert_params],
         [local_non_expert_param_hashes, local_expert_param_hashes],
-        [
-            parallel_state.get_data_parallel_group(with_gtp_remat=False),
-            parallel_state.get_expert_data_parallel_group(with_gtp_remat=False),
-        ],
+        [dp_group, expt_dp_group],
     ):
         # Collect per-parameter hashes across all ranks in group.
         assert len(params) == len(local_param_hashes)

megatron/training/training.py

@@ -4443,6 +4443,23 @@ def _run_gpu_sniff_test(tag, span_name='megatron.train.sniff_test'):
     print_datetime(f'finished GPU sniff test ({tag})')
 
 
+def _param_hashes_match_across_dp_replicas(model):
+    """Return whether every data-parallel replica of the model holds the same parameters.
+
+    The replicas are those of the model's process-group collection; a model without one is
+    checked across the replicas of the global parallel state.
+    """
+    model_pg_collection = get_attr_wrapped_model(model[0], "pg_collection")
+    if model_pg_collection is None:
+        model_pg_collection = ProcessGroupCollection.use_mpu_process_groups()
+    return check_param_hashes_across_dp_replicas(
+        model,
+        cross_check=True,
+        dp_group=model_pg_collection.dp,
+        expt_dp_group=model_pg_collection.expt_dp,
+    )
+
+
 def post_training_step_callbacks(
     model,
     optimizer,
@@ -4475,8 +4492,8 @@ def post_training_step_callbacks(
     ):
         if should_disable_forward_pre_hook(args, optimizer):
             disable_forward_pre_hook(model, optimizer=optimizer)
-        assert check_param_hashes_across_dp_replicas(
-            model, cross_check=True
+        assert _param_hashes_match_across_dp_replicas(
+            model
         ), "Parameter hashes not matching across DP replicas"
         torch.distributed.barrier()
         print_rank_0(f">>> Weight hashes match after {iteration} iterations...")
@@ -4991,8 +5008,8 @@ def trace_handler(p):
     # span (parented to megatron.startup, still open here).
     if cfg.train.check_weight_hash_across_dp_replicas_interval is not None:
         with _otel_managed_span('job', 'megatron.startup.weight_hash_check', is_goodput_span=True):
-            assert check_param_hashes_across_dp_replicas(
-                model, cross_check=True
+            assert _param_hashes_match_across_dp_replicas(
+                model
             ), "Parameter hashes not matching across DP replicas"
             torch.distributed.barrier()
         print_rank_0(f">>> Weight hashes match after {iteration} iterations...")

tests/unit_tests/test_utils.py

@@ -1,5 +1,6 @@
 # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
 
+import contextlib
 import os
 import time
 import urllib.request as req
@@ -13,9 +14,10 @@
 
 import megatron.core.utils as util
 import megatron.training.utils as training_util
-from megatron.core import config
+from megatron.core import config, parallel_state
 from megatron.core._rank_utils import safe_get_rank, safe_get_world_size
 from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig
+from megatron.core.hyper_comm_grid import HyperCommGrid
 from megatron.core.models.gpt.gpt_layer_specs import (
     get_gpt_layer_with_transformer_engine_submodules,
 )
@@ -326,6 +328,118 @@ def test_cross_check_param_hashes_across_dp_replicas():
     _deinit_distributed()
 
 
+@contextlib.contextmanager
+def _forbid_global_process_groups():
+    """Make every parallel_state group, rank and world-size accessor raise when called."""
+
+    def _raise(*args, **kwargs):
+        raise RuntimeError("read a process group of the global parallel state")
+
+    with contextlib.ExitStack() as stack:
+        for name in dir(parallel_state):
+            if name.startswith("get_") and name.endswith(
+                ("_group", "_groups", "_gloo", "_rank", "_ranks", "_world_size")
+            ):
+                stack.enter_context(patch.object(parallel_state, name, _raise))
+        yield
+
+
+def _dense_and_expert_model(dense_value, expert_value):
+    """Return a model chunk with one dense and one expert (``allreduce=False``) parameter."""
+    model = torch.nn.ModuleDict(
+        {
+            "dense": torch.nn.Linear(8, 8, bias=False, device="cuda"),
+            "expert": torch.nn.Linear(8, 8, bias=False, device="cuda"),
+        }
+    )
+    model.dense.weight.data.fill_(dense_value)
+    model.expert.weight.data.fill_(expert_value)
+    model.expert.weight.allreduce = False
+    return model
+
+
+@pytest.mark.skipif(
+    Utils.world_size < 4 or Utils.world_size % 2 != 0,
+    reason="needs an even number of ranks, at least 4",
+)
+class TestCheckParamHashesWithModelGroups:
+    """check_param_hashes_across_dp_replicas compares hashes across the groups it is given.
+
+    The global grid has one data-parallel group and one expert data-parallel group, each of
+    all ranks. The given groups are laid out differently: dense replicas are pairs of adjacent
+    ranks, and expert replicas are the ranks of equal parity. Parameters are equal within each
+    given group and differ across groups, so comparing across a global group fails.
+    """
+
+    def setup_method(self, method):
+        Utils.initialize_model_parallel()
+        self.grid = HyperCommGrid([2, torch.distributed.get_world_size() // 2], ["dp", "expt_dp"])
+        self.dp_group = self.grid.create_pg("dp")
+        self.expt_dp_group = self.grid.create_pg("expt_dp")
+        self.dp_ranks = torch.distributed.get_process_group_ranks(self.dp_group)
+        self.expt_dp_ranks = torch.distributed.get_process_group_ranks(self.expt_dp_group)
+
+    def teardown_method(self, method):
+        self.grid.destroy()
+        Utils.destroy_model_parallel()
+
+    def _model(self):
+        return _dense_and_expert_model(float(self.dp_ranks[0]), float(self.expt_dp_ranks[0]))
+
+    @pytest.mark.parametrize("cross_check", [False, True])
+    def test_compares_across_the_given_groups(self, cross_check):
+        model = self._model()
+
+        with _forbid_global_process_groups():
+            assert util.check_param_hashes_across_dp_replicas(
+                [model], cross_check, dp_group=self.dp_group, expt_dp_group=self.expt_dp_group
+            )
+        # Without groups, the check compares across the global groups, which span ranks whose
+        # parameters differ.
+        assert not util.check_param_hashes_across_dp_replicas([model], cross_check=True)
+
+    @pytest.mark.parametrize("cross_check", [False, True])
+    def test_detects_a_mismatch_within_the_given_groups(self, cross_check):
+        rank = torch.distributed.get_rank()
+        model = self._model()
+        # Rank 1 changes a dense parameter and rank 2 an expert parameter. Neither is the first
+        # rank of the group across which that parameter is compared.
+        if rank == 1:
+            model.dense.weight.data.add_(100.0)
+        if rank == 2:
+            model.expert.weight.data.add_(100.0)
+
+        with _forbid_global_process_groups():
+            hashes_match = util.check_param_hashes_across_dp_replicas(
+                [model], cross_check, dp_group=self.dp_group, expt_dp_group=self.expt_dp_group
+            )
+
+        if cross_check:
+            # Every rank of a group that contains a changed parameter sees the mismatch.
+            expected = 1 not in self.dp_ranks and 2 not in self.expt_dp_ranks
+        else:
+            # Each rank compares with the first rank of its group, whose parameters are unchanged.
+            expected = rank not in (1, 2)
+        assert hashes_match == expected
+
+    def test_requires_the_group_of_each_parameter_kind(self):
+        model = self._model()
+        dense_model = torch.nn.Linear(8, 8, bias=False, device="cuda")
+        dense_model.weight.data.fill_(float(self.dp_ranks[0]))
+
+        with _forbid_global_process_groups():
+            with pytest.raises(ValueError, match="has expert parameters"):
+                util.check_param_hashes_across_dp_replicas([model], dp_group=self.dp_group)
+            with pytest.raises(ValueError, match="has non-expert parameters"):
+                util.check_param_hashes_across_dp_replicas(
+                    [model], expt_dp_group=self.expt_dp_group
+                )
+            # A model without expert parameters needs no expert group.
+            assert util.check_param_hashes_across_dp_replicas(
+                [dense_model], True, dp_group=self.dp_group
+            )
+
+
 @pytest.mark.parametrize("use_distributed_optimizer", [False, True])
 @pytest.mark.flaky
 @pytest.mark.flaky_in_dev

tests/unit_tests/training/test_param_hash_check_groups.py

@@ -0,0 +1,96 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+"""The training loop compares parameter hashes across the model's own data-parallel replicas."""
+
+from types import SimpleNamespace
+from unittest import mock
+
+from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.training import training as training_mod
+
+
+def _single_call_kwargs(check, model):
+    """Return the keyword arguments of the one call to the mocked check, which got model."""
+    assert check.call_count == 1
+    (checked_model,), kwargs = check.call_args
+    assert checked_model is model
+    return kwargs
+
+
+def _hash_check_kwargs(model):
+    """Run the training loop's hash check on model; return the arguments the core check got."""
+    check = mock.Mock(return_value=True)
+    with mock.patch.object(training_mod, "check_param_hashes_across_dp_replicas", check):
+        assert training_mod._param_hashes_match_across_dp_replicas(model)
+    return _single_call_kwargs(check, model)
+
+
+def test_compares_across_the_model_dp_groups():
+    pg_collection = ProcessGroupCollection(dp=object(), expt_dp=object())
+
+    kwargs = _hash_check_kwargs([SimpleNamespace(pg_collection=pg_collection)])
+
+    assert kwargs == {
+        "cross_check": True,
+        "dp_group": pg_collection.dp,
+        "expt_dp_group": pg_collection.expt_dp,
+    }
+
+
+def test_collection_without_expert_groups_passes_no_expert_group():
+    pg_collection = ProcessGroupCollection(dp=object())
+
+    kwargs = _hash_check_kwargs([SimpleNamespace(pg_collection=pg_collection)])
+
+    assert kwargs["dp_group"] is pg_collection.dp
+    assert kwargs["expt_dp_group"] is None
+
+
+def test_model_without_collection_uses_the_global_groups():
+    global_pg_collection = ProcessGroupCollection(dp=object(), expt_dp=object())
+
+    with mock.patch.object(
+        ProcessGroupCollection, "use_mpu_process_groups", return_value=global_pg_collection
+    ):
+        kwargs = _hash_check_kwargs([SimpleNamespace(pg_collection=None)])
+
+    assert kwargs["dp_group"] is global_pg_collection.dp
+    assert kwargs["expt_dp_group"] is global_pg_collection.expt_dp
+
+
+def test_post_training_step_callbacks_check_the_model_dp_groups():
+    pg_collection = ProcessGroupCollection(dp=object(), expt_dp=object())
+    model = [SimpleNamespace(pg_collection=pg_collection)]
+    # The training settings live in args before the TrainingConfig migration (#8149) and in
+    # the run config's train section after it, so the test sets both.
+    train = dict(
+        train_sync_interval=0,
+        check_weight_hash_across_dp_replicas_interval=2,
+        gpu_sniff_test_interval=None,
+        manual_gc=False,
+        manual_gc_interval=0,
+    )
+    args = SimpleNamespace(log_straggler=False, adlr_autoresume=False, **train)
+    cfg = SimpleNamespace(
+        train=SimpleNamespace(**train),
+        logger=SimpleNamespace(log_interval=1),
+        profiling=SimpleNamespace(use_nsys_profiler=False, use_pytorch_profiler=False),
+    )
+    check = mock.Mock(return_value=True)
+
+    with (
+        mock.patch.object(training_mod, "get_args", return_value=args),
+        mock.patch.object(training_mod, "get_run_config", return_value=cfg),
+        mock.patch.object(training_mod, "should_disable_forward_pre_hook", return_value=False),
+        mock.patch.object(training_mod, "check_param_hashes_across_dp_replicas", check),
+        mock.patch.object(training_mod.torch.distributed, "barrier"),
+        mock.patch.object(training_mod, "print_rank_0"),
+    ):
+        for iteration in (1, 2):
+            training_mod.post_training_step_callbacks(model, None, None, iteration, None, 0)
+
+    assert _single_call_kwargs(check, model) == {
+        "cross_check": True,
+        "dp_group": pg_collection.dp,
+        "expt_dp_group": pg_collection.expt_dp,
+    }