[sync][training] fix: honor checkpoint sharding validation policy

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>

Author: yaoyu-33

src/megatron/bridge/training/checkpointing.py

@@ -119,6 +119,11 @@
 _get_legacy_grad_dtypes = getattr(distrib_optimizer, "get_legacy_grad_dtypes", None)
 _CHECKPOINT_FORMAT_VERSION = 3.1 if _get_legacy_grad_dtypes is not None else 3.0
 
+# The supported dev pin exposes a load-only policy and has no save-side opt-out.
+_FULLY_PARALLEL_SAVE_HAS_VALIDATION_POLICY = (
+    "validate_access_integrity" in inspect.signature(FullyParallelSaveStrategyWrapper).parameters
+)
+
 
 _, HAVE_RESIL = safe_import("nvidia_resiliency_ext.checkpointing")
 
@@ -1532,11 +1537,19 @@ def save_checkpoint(
                             logger.debug("Failed to plug in the read metadata from the load strategy...")
 
                 if ckpt_cfg.fully_parallel_save:
+                    validation_kwargs = (
+                        {"validate_access_integrity": ckpt_cfg.ckpt_load_validate_sharding_integrity}
+                        if _FULLY_PARALLEL_SAVE_HAS_VALIDATION_POLICY
+                        else {}
+                    )
                     save_strategy = FullyParallelSaveStrategyWrapper(
                         save_strategy,
                         _checkpoint_dp_cp_group(pg_collection),
                         ckpt_cfg.ckpt_assume_constant_structure,
+                        **validation_kwargs,
                     )
+            if _FULLY_PARALLEL_SAVE_HAS_VALIDATION_POLICY and not ckpt_cfg.ckpt_load_validate_sharding_integrity:
+                validate_sharding_integrity = False
             # MegatronMIMO + torch_dist can hit known access-pattern validation failures
             # for nested DDP language model tensors in PP>1 runs when
             # fully_parallel_save is disabled. Keep validation enabled otherwise.
@@ -3839,7 +3852,7 @@ def _load_global_dist_base_checkpoint(
         load_strategy = FullyParallelLoadStrategyWrapper(load_strategy, _checkpoint_dp_cp_group(pg_collection))
     if checkpointing_context is not None:
         checkpointing_context["load_strategy"] = load_strategy
-    validate_sharding_integrity = True
+    validate_sharding_integrity = ckpt_cfg.ckpt_load_validate_sharding_integrity
     if is_megatron_mimo and ckpt_cfg.ckpt_format == "torch_dist" and not ckpt_cfg.fully_parallel_save:
         validate_sharding_integrity = False
     if cfg is not None:

tests/unit_tests/training/test_checkpointing.py

@@ -56,6 +56,7 @@
     _has_global_non_persistent_checkpoint,
     _load_base_checkpoint,
     _load_checkpoint_from_path,
+    _load_global_dist_base_checkpoint,
     _load_hf_pretrained_checkpoint,
     _load_model_state_dict,
     _load_non_persistent_base_checkpoint,
@@ -846,7 +847,22 @@ def test_metadata_failure_does_not_publish_incomplete_checkpoint(self, tmp_path,
         assert torch.load(latest_train_state, weights_only=True)["step"].item() == 500
         assert legacy_tracker.read_text() == "500"
 
+    @pytest.mark.parametrize("save_validation_supported", [True, False])
     @pytest.mark.parametrize("save_rng", [True, False])
+    @pytest.mark.parametrize("validate_sharding", [True, False])
+    @pytest.mark.parametrize(
+        "fully_parallel,cached_strategy,constant_structure,is_mimo",
+        [
+            (False, False, False, False),
+            (True, False, False, False),
+            (False, True, True, False),
+            (True, True, True, False),
+            (True, True, False, False),
+            (False, False, False, True),
+            (True, False, False, True),
+        ],
+    )
+    @patch("megatron.bridge.training.checkpointing.FullyParallelSaveStrategyWrapper")
     @patch("megatron.bridge.training.checkpointing.wandb_utils")
     @patch("megatron.bridge.training.checkpointing.is_last_rank")
     @patch("builtins.open", new_callable=mock_open)
@@ -893,10 +909,29 @@ def test_save_checkpoint_global(
         mock_file_open,
         mock_is_last_rank,
         mock_wandb,
+        mock_parallel_strategy,
         save_checkpoint_fixtures,
+        monkeypatch,
+        save_validation_supported,
         save_rng,
+        validate_sharding,
+        fully_parallel,
+        cached_strategy,
+        constant_structure,
+        is_mimo,
     ):
         """Test saving a global checkpoint."""
+        monkeypatch.setattr(
+            "megatron.bridge.training.checkpointing._FULLY_PARALLEL_SAVE_HAS_VALIDATION_POLICY",
+            save_validation_supported,
+            raising=False,
+        )
+        if not save_validation_supported:
+            # Match the supported dev pin's constructor: a new keyword must fail this test.
+            def legacy_wrapper(strategy, parallelization_group, do_cache_distribution):
+                return Mock()
+
+            mock_parallel_strategy.side_effect = legacy_wrapper
         # Setup mocks
         mock_dist_init.return_value = True
         mock_get_dist_rank.return_value = 0
@@ -929,6 +964,12 @@ def test_save_checkpoint_global(
         save_checkpoint_fixtures["mock_state"].wandb_logger = Mock()
         save_checkpoint_fixtures["mock_state"].cfg.checkpoint.most_recent_k = -1
         save_checkpoint_fixtures["mock_state"].cfg.checkpoint.save_rng = save_rng
+        save_checkpoint_fixtures["mock_state"].cfg.checkpoint.ckpt_load_validate_sharding_integrity = validate_sharding
+        save_checkpoint_fixtures["mock_state"].cfg.checkpoint.fully_parallel_save = fully_parallel
+        save_checkpoint_fixtures["mock_state"].cfg.checkpoint.ckpt_assume_constant_structure = constant_structure
+        if not is_mimo:
+            del save_checkpoint_fixtures["mock_model"][0].mimo_config
+        checkpoint_context = {"save_strategy": Mock()} if cached_strategy else {}
 
         call_order = Mock()
         call_order.attach_mock(mock_dist_ckpt.save, "model")
@@ -941,7 +982,7 @@ def test_save_checkpoint_global(
             save_checkpoint_fixtures["mock_optimizer"],
             save_checkpoint_fixtures["mock_scheduler"],
             1000000,
-            checkpointing_context={},
+            checkpointing_context=checkpoint_context,
             non_persistent_ckpt=False,
         )
 
@@ -951,6 +992,19 @@ def test_save_checkpoint_global(
         mock_ft.on_checkpointing_start.assert_called_once()
         mock_gen_state.assert_called_once()
         mock_dist_ckpt.save.assert_called_once()
+        expected_validation = (validate_sharding or not save_validation_supported) and not (
+            (cached_strategy and constant_structure) or (is_mimo and not fully_parallel)
+        )
+        assert mock_dist_ckpt.save.call_args.kwargs["validate_access_integrity"] is expected_validation
+        if fully_parallel and not cached_strategy:
+            if save_validation_supported:
+                assert (
+                    mock_parallel_strategy.call_args.kwargs.get("validate_access_integrity", True) is validate_sharding
+                )
+            else:
+                assert "validate_access_integrity" not in mock_parallel_strategy.call_args.kwargs
+        else:
+            mock_parallel_strategy.assert_not_called()
         if save_rng:
             mock_get_rng.assert_called_once()
         else:
@@ -6836,3 +6890,30 @@ def test_checkpoint_minor_version_changes_are_compatible(monkeypatch, versions):
     assert checkpointing.get_checkpoint_version() == versions[-1]
     with pytest.raises(AssertionError, match="checkpoint versions do not match"):
         checkpointing.set_checkpoint_version(2.0)
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("validate_sharding", [True, False])
+@pytest.mark.parametrize("fully_parallel", [True, False])
+@pytest.mark.parametrize("is_mimo", [True, False])
+def test_global_checkpoint_load_honors_sharding_validation(
+    validate_sharding: bool, fully_parallel: bool, is_mimo: bool
+) -> None:
+    """Public checkpoint policy must reach the Core loader."""
+    cfg = CheckpointConfig(
+        ckpt_format="torch_dist",
+        ckpt_load_validate_sharding_integrity=validate_sharding,
+        fully_parallel_load=fully_parallel,
+        fully_parallel_save=fully_parallel,
+    )
+    pg_collection = Mock()
+    with (
+        patch("megatron.bridge.training.checkpointing.TorchDistLoadShardedStrategy"),
+        patch("megatron.bridge.training.checkpointing.FullyParallelLoadStrategyWrapper"),
+        patch("megatron.bridge.training.checkpointing.dist_checkpointing.load", return_value={}) as load,
+    ):
+        _load_global_dist_base_checkpoint(
+            "checkpoints", cfg, False, {"model": {}}, 1, False, is_megatron_mimo=is_mimo, pg_collection=pg_collection
+        )
+    expected_validation = validate_sharding and not (is_mimo and not fully_parallel)
+    assert load.call_args.kwargs["validate_access_integrity"] is expected_validation