From 128b1db76bbfbf60082d79afa1c14f772344b476 Mon Sep 17 00:00:00 2001 From: Yan Xu Date: Fri, 9 Oct 2026 17:06:58 -0700 Subject: [PATCH 1/3] Let check_param_hashes_across_dp_replicas take the model's DP groups check_param_hashes_across_dp_replicas all-gathered parameter hashes over the data-parallel groups of the global parallel state. A model whose process-group collection differs from the global grid then compares its hashes with ranks that hold other shards or another model: the check reports a mismatch for replicas that agree, or hangs when the ranks hold different numbers of parameters. Add keyword-only dp_group and expt_dp_group arguments for the model's replicate groups, pg_collection.dp and pg_collection.expt_dp. These are the fields that match the global accessors called with with_gtp_remat=False. If neither is passed, the function keeps using the global groups as a commented migration fallback. If either is passed, no global group is read. A rank that holds non-expert parameters without dp_group, or expert parameters without expt_dp_group, raises ValueError before the first collective, so a dense model may pass dp_group alone. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Yan Xu --- megatron/core/utils.py | 43 ++++++++++-- tests/unit_tests/test_utils.py | 116 ++++++++++++++++++++++++++++++++- 2 files changed, 153 insertions(+), 6 deletions(-) diff --git a/megatron/core/utils.py b/megatron/core/utils.py index 4d47cb24b37..37f3fb12aff 100644 --- a/megatron/core/utils.py +++ b/megatron/core/utils.py @@ -972,7 +972,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. @@ -981,15 +985,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 @@ -1011,16 +1034,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) diff --git a/tests/unit_tests/test_utils.py b/tests/unit_tests/test_utils.py index 24836fcb435..2d0b44e940d 100644 --- a/tests/unit_tests/test_utils.py +++ b/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 From 0ab8e49443389be5719f34f0bcba046f48970ac7 Mon Sep 17 00:00:00 2001 From: Yan Xu Date: Fri, 9 Oct 2026 17:06:59 -0700 Subject: [PATCH 2/3] Check weight hashes across the model's DP groups in training The training loop checked parameter hashes at startup and every --check-weight-hash-across-dp-replicas-interval iterations across the global data-parallel groups. Pass the data-parallel groups of the model's process-group collection instead, resolved as train_step does: the collection of the model chunks, or the global one for a model that carries none. On the standard grid the model's collection holds the global groups, so the check compares across the same ranks as before. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Yan Xu --- megatron/training/training.py | 25 ++++- .../training/test_param_hash_check_groups.py | 93 +++++++++++++++++++ 2 files changed, 114 insertions(+), 4 deletions(-) create mode 100644 tests/unit_tests/training/test_param_hash_check_groups.py diff --git a/megatron/training/training.py b/megatron/training/training.py index 27e2d5eff59..36592172f1b 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -4423,6 +4423,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, @@ -4455,8 +4472,8 @@ def post_training_step_callbacks( ): if should_disable_forward_pre_hook(args): 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...") @@ -5005,8 +5022,8 @@ def trace_handler(p): # span (parented to megatron.startup, still open here). if args.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...") diff --git a/tests/unit_tests/training/test_param_hash_check_groups.py b/tests/unit_tests/training/test_param_hash_check_groups.py new file mode 100644 index 00000000000..b892f3a423e --- /dev/null +++ b/tests/unit_tests/training/test_param_hash_check_groups.py @@ -0,0 +1,93 @@ +# 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)] + args = SimpleNamespace( + train_sync_interval=0, + log_straggler=False, + check_weight_hash_across_dp_replicas_interval=2, + adlr_autoresume=False, + gpu_sniff_test_interval=None, + manual_gc=False, + ) + cfg = SimpleNamespace( + 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, + } From 8f2885e7c51f5785a4c75395e94c45e847d0349a Mon Sep 17 00:00:00 2001 From: Yan Xu Date: Tue, 27 Oct 2026 16:19:32 -0700 Subject: [PATCH 3/3] Set the run config's train section in the callback hash-check test The TrainingConfig migration (#8149) moved train_sync_interval, the weight hash check interval, the GPU sniff test interval and the manual GC settings from args to the run config's train section. The test only set them in args, so after a rebase past #8149, post_training_step_callbacks failed on cfg.train. Set the same values in both, which passes before and after the migration. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Yan Xu --- .../unit_tests/training/test_param_hash_check_groups.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tests/unit_tests/training/test_param_hash_check_groups.py b/tests/unit_tests/training/test_param_hash_check_groups.py index b892f3a423e..cd8468e08aa 100644 --- a/tests/unit_tests/training/test_param_hash_check_groups.py +++ b/tests/unit_tests/training/test_param_hash_check_groups.py @@ -61,15 +61,18 @@ def test_model_without_collection_uses_the_global_groups(): 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)] - args = SimpleNamespace( + # 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, - log_straggler=False, check_weight_hash_across_dp_replicas_interval=2, - adlr_autoresume=False, 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), )