feat(ckpt): expose validate_access_integrity knob on dist-ckpt load

megatron/training/checkpointing.py

@@ -1247,6 +1247,7 @@ def _load_global_dist_base_checkpoint(
         sharded_state_dict,
         checkpoint_name,
         load_strategy,
+        validate_access_integrity=args.ckpt_load_validate_sharding_integrity,
         strict=args.dist_ckpt_strictness,
     )
     return state_dict, checkpoint_name, release, CheckpointType.GLOBAL

megatron/training/config/training_config.py

@@ -529,6 +529,11 @@ class CheckpointConfig:
     ckpt_assume_constant_structure: bool = False
     """Assume the checkpoint structure is constant across saves to enable optimizations."""
 
+    ckpt_load_validate_sharding_integrity: bool = True
+    """Whether to validate sharding access integrity when loading a distributed checkpoint.
+    When True (default), each tensor shard is checked to be accessed exactly once as main
+    replica by some rank. Disabling skips this validation"""
+
     strict_fsdp_dtensor_load: bool = True
     """Whether to enforce strict loading for FSDP DTensor checkpoints. When False, allows partial loading."""