send pg group for distributed checkpoint validation

megatron/core/dist_checkpointing/serialization.py

@@ -73,6 +73,7 @@ def load(
     validate_access_integrity: bool = True,
     strict: Union[str, StrictHandling] = StrictHandling.ASSUME_OK_UNEXPECTED,
     verify_integrity: bool = False,
+    process_group: Optional[torch.distributed.ProcessGroup] = None,
 ) -> Union[StateDict, Tuple[StateDict, Set[str], Set[str]]]:
     """Loading entrypoint.
 
@@ -108,6 +109,8 @@ def load(
             and compares against the SHA-256 manifest. Raises `CheckpointingException` on any
             mismatch. Requires that the checkpoint was previously saved with
             `verify_integrity=True`.
+        process_group (ProcessGroup, optional): ranks that collectively describe
+            one complete sharded state dict. Defaults to the global process group.
 
     Returns:
         StateDict or Tuple[StateDict, Set[str], Set[str]]: in most cases only
@@ -163,7 +166,9 @@ def load(
             k: v for k, v in ckpt_sharded_metadata.items() if v.key != 'common_state'
         }
     if validate_access_integrity or StrictHandling.requires_global_app_metadata(strict):
-        local_metadata, global_metadata = determine_global_metadata(sharded_state_dict)
+        local_metadata, global_metadata = determine_global_metadata(
+            sharded_state_dict, process_group=process_group
+        )
 
     sharded_state_dict, missing_keys, unexpected_keys = validate_integrity_and_strict_load(
         sharded_state_dict,

megatron/core/dist_checkpointing/validation.py

@@ -477,18 +477,21 @@ def _validate_objects_for_key(sharded_objects: List[ShardedObject]) -> List[Chec
 
 def determine_global_metadata(
     sharded_state_dict: ShardedStateDict,
+    process_group: Optional[torch.distributed.ProcessGroup] = None,
 ) -> Tuple[_LocalMetadata, _GlobalMetadata]:
     """Exchanges local metadata with `all_gather_object` to determine global metadata.
 
     Args:
         sharded_state_dict (ShardedStateDict): local sharded state dict
+        process_group (ProcessGroup, optional): ranks whose metadata forms one
+            complete checkpoint view. Defaults to the global process group.
 
     Returns:
         Tuple[_LocalMetadata, _GlobalMetadata]: local and global ShardedBase objects with stripped data
     """
     local_metadata = [ten.without_data() for ten in nested_values(sharded_state_dict)]
-    global_metadata = [None] * torch.distributed.get_world_size()
-    torch.distributed.all_gather_object(global_metadata, local_metadata)
+    global_metadata = [None] * torch.distributed.get_world_size(group=process_group)
+    torch.distributed.all_gather_object(global_metadata, local_metadata, group=process_group)
     return local_metadata, global_metadata  # type: ignore[return-value]
 
 

tests/unit_tests/dist_checkpointing/test_validation.py

@@ -0,0 +1,27 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+from unittest.mock import Mock
+
+import torch
+
+from megatron.core.dist_checkpointing.validation import determine_global_metadata
+
+
+def test_determine_global_metadata_uses_explicit_process_group(monkeypatch):
+    process_group = Mock()
+    metadata = Mock()
+    shard = Mock()
+    shard.without_data.return_value = metadata
+    get_world_size = Mock(return_value=2)
+    all_gather_object = Mock()
+    monkeypatch.setattr(torch.distributed, "get_world_size", get_world_size)
+    monkeypatch.setattr(torch.distributed, "all_gather_object", all_gather_object)
+
+    local_metadata, global_metadata = determine_global_metadata(
+        {"model": shard}, process_group=process_group
+    )
+
+    assert local_metadata == [metadata]
+    assert global_metadata == [None, None]
+    get_world_size.assert_called_once_with(group=process_group)
+    all_gather_object.assert_called_once_with(global_metadata, local_metadata, group=process_group)