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."""
 

tests/unit_tests/dist_checkpointing/test_pipeline_parallel_layout.py

@@ -140,6 +140,7 @@ def create_args():
     args.vocab_file = None
     args.add_position_embedding = False
     args.ckpt_assume_constant_structure = True
+    args.ckpt_load_validate_sharding_integrity = True
     args.dist_ckpt_strictness = "assume_ok_unexpected"
     args.fp16 = False
     args.bf16 = True

tests/unit_tests/dist_checkpointing/utils.py

@@ -150,6 +150,7 @@ def init_checkpointing_mock_args(args, ckpt_dir, fully_parallel=False):
     args.no_save_optim = False
     args.no_save_rng = False
     args.ckpt_assume_constant_structure = False
+    args.ckpt_load_validate_sharding_integrity = True
     args.log_progress = False
     args.auto_detect_ckpt_format = False
     args.exit_on_missing_checkpoint = False

tests/unit_tests/pipeline_parallel/test_pipeline_layout.py

@@ -135,6 +135,7 @@ def create_args():
     args.vocab_file = None
     args.add_position_embedding = False
     args.ckpt_assume_constant_structure = False
+    args.ckpt_load_validate_sharding_integrity = True
     args.dist_ckpt_strictness = "assume_ok_unexpected"
     args.fp16 = False
     args.bf16 = True

tests/unit_tests/test_checkpointing.py

@@ -140,6 +140,7 @@ def create_ckpt_load_args(create_args):
     args.ckpt_assume_constant_structure = False
     args.ckpt_fully_parallel_save = False
     args.ckpt_fully_parallel_load = False
+    args.ckpt_load_validate_sharding_integrity = True
     args.dist_ckpt_strictness = 'assume_ok_unexpected'
     args.use_megatron_fsdp = False
     args.strict_fsdp_dtensor_load = True