fix: harden checkpoint config target validation

src/megatron/bridge/training/model_load_save.py

@@ -235,7 +235,7 @@ def load_model_config(
         read_run_config,
     )
     from megatron.bridge.training.mlm_compat.arguments import _load_args_from_checkpoint, _transformer_config_from_args
-    from megatron.bridge.utils.instantiate_utils import instantiate
+    from megatron.bridge.utils.instantiate_utils import instantiate, validate_config_targets
 
     run_config_filename = get_checkpoint_run_config_filename(checkpoint_path)
 
@@ -259,6 +259,7 @@ def load_model_config(
 
     if mbridge_ckpt:
         if "_builder_" in run_config["model"]:
+            validate_config_targets(run_config["model"], target_keys=("_target_", "_builder_"))
             model_cfg = ModelConfig.from_dict(run_config["model"])
         else:
             model_cfg = instantiate(run_config["model"])

src/megatron/bridge/utils/instantiate_utils.py

@@ -15,6 +15,8 @@
 # Patch for https://github.com/facebookresearch/hydra/blob/main/hydra/_internal/instantiate/_instantiate2.py
 # until https://github.com/facebookresearch/hydra/issues/2140 is resolved
 
+from typing import Any, NamedTuple
+
 from megatron.training.config.instantiate_utils import (
     InstantiationException,
     InstantiationMode,  # noqa: F401  (re-exported for tests / external callers)
@@ -28,10 +30,16 @@
     _locate,  # noqa: F401  (re-exported for tests / external callers)
     _prepare_input_dict_or_list,  # noqa: F401  (re-exported for tests / external callers)
     _resolve_target,  # noqa: F401  (re-exported for tests / external callers)
-    instantiate,  # noqa: F401  (re-exported for tests / external callers)
-    instantiate_node,  # noqa: F401  (re-exported for tests / external callers)
     target_allowlist,
 )
+from megatron.training.config.instantiate_utils import (
+    instantiate as _mlm_instantiate,
+)
+from megatron.training.config.instantiate_utils import (
+    instantiate_node as _mlm_instantiate_node,
+)
+from omegaconf import OmegaConf
+from omegaconf._utils import is_structured_config
 
 
 _ALLOWED_TARGET_PREFIXES: set[str] = {
@@ -44,6 +52,25 @@
 }
 
 
+class _TargetAllowlistSnapshot(NamedTuple):
+    """Immutable view of the target allowlist used for one config preflight."""
+
+    enabled: bool
+    allowed_prefixes: tuple[str, ...]
+    allowed_exact: frozenset[str]
+
+
+_BLOCKED_TARGETS: frozenset[str] = frozenset(
+    {
+        "megatron.bridge.utils.instantiate_utils.register_allowed_target_prefix",
+    }
+)
+_BLOCKED_TARGET_PREFIXES: tuple[str, ...] = (
+    "megatron.bridge.utils.instantiate_utils.target_allowlist.",
+    "megatron.training.config.instantiate_utils.target_allowlist.",
+)
+
+
 # Mirror Bridge's allowlist into the MLM `target_allowlist` singleton, which is
 # the source of truth consulted by `_validate_target_prefix` below. MLM's
 # default prefixes are narrower (megatron.training./megatron.core./torch./
@@ -62,6 +89,122 @@ def _seed_allowlist() -> None:
 _seed_allowlist()
 
 
+def _snapshot_allowlist() -> _TargetAllowlistSnapshot:
+    """Capture the allowlist state before any target in a config can run."""
+    return _TargetAllowlistSnapshot(
+        enabled=target_allowlist.enabled,
+        allowed_prefixes=target_allowlist.allowed_prefixes,
+        allowed_exact=target_allowlist.allowed_exact,
+    )
+
+
+def _is_allowed_by_snapshot(target: str, snapshot: _TargetAllowlistSnapshot) -> bool:
+    """Check a target against the pre-instantiation allowlist snapshot."""
+    if not snapshot.enabled:
+        return True
+    if target in snapshot.allowed_exact:
+        return True
+    return any(target.startswith(prefix) for prefix in snapshot.allowed_prefixes)
+
+
+def _target_is_blocked(target: str) -> bool:
+    """Reject config targets that can mutate the allowlist during traversal."""
+    return target in _BLOCKED_TARGETS or any(target.startswith(prefix) for prefix in _BLOCKED_TARGET_PREFIXES)
+
+
+def _raise_disallowed_target(target: str, full_key: str, snapshot: _TargetAllowlistSnapshot) -> None:
+    """Raise an allowlist error using the stable snapshot for diagnostics."""
+    raise InstantiationException(
+        f"Target '{target}' is not in the allowlist for _target_ instantiation.\n"
+        f"Allowed module prefixes: {', '.join(snapshot.allowed_prefixes)}\n"
+        f"Allowed exact targets: {', '.join(sorted(snapshot.allowed_exact))}"
+        + (f"\nfull_key: {full_key}" if full_key else "")
+    )
+
+
+def _validate_target_for_instantiate(target: Any, full_key: str, snapshot: _TargetAllowlistSnapshot) -> None:
+    """Validate one target before import or recursive instantiation can occur."""
+    target = _convert_target_to_string(target)
+    if not isinstance(target, str):
+        return
+    if _target_is_blocked(target):
+        raise InstantiationException(
+            f"Target '{target}' is not allowed in config instantiation because it can modify the target allowlist."
+            + (f"\nfull_key: {full_key}" if full_key else "")
+        )
+    if not _is_allowed_by_snapshot(target, snapshot):
+        _raise_disallowed_target(target, full_key, snapshot)
+
+
+def _preflight_targets(
+    node: Any,
+    snapshot: _TargetAllowlistSnapshot,
+    full_key: str = "",
+    target_keys: tuple[str, ...] = (_Keys.TARGET.value,),
+) -> None:
+    """Validate every target in a config tree against the same allowlist snapshot."""
+    if node is None:
+        return
+
+    if isinstance(node, (dict, list)):
+        node = _prepare_input_dict_or_list(node)
+        node = OmegaConf.structured(node, flags={"allow_objects": True})
+    elif is_structured_config(node):
+        node = OmegaConf.structured(node, flags={"allow_objects": True})
+
+    if OmegaConf.is_dict(node):
+        for target_key in target_keys:
+            if target_key in node:
+                target_full_key = target_key if not full_key else f"{full_key}.{target_key}"
+                _validate_target_for_instantiate(node.get(target_key), target_full_key, snapshot)
+        for key in node.keys():
+            if key in (*target_keys, _Keys.PARTIAL, _Keys.CALL, _Keys.NAME):
+                continue
+            child_key = str(key) if not full_key else f"{full_key}.{key}"
+            _preflight_targets(node[key], snapshot, child_key, target_keys)
+    elif OmegaConf.is_list(node):
+        for idx, value in enumerate(node._iter_ex(resolve=True)):
+            child_key = f"{full_key}[{idx}]" if full_key else f"[{idx}]"
+            _preflight_targets(value, snapshot, child_key, target_keys)
+
+
+def validate_config_targets(config: Any, target_keys: tuple[str, ...] = (_Keys.TARGET.value,)) -> None:
+    """Validate code-selection fields in config metadata without importing them.
+
+    Args:
+        config: Config tree to validate.
+        target_keys: Dictionary keys whose string values select Python objects.
+            Checkpoint metadata uses both ``_target_`` and ``_builder_``.
+    """
+    snapshot = _snapshot_allowlist()
+    _preflight_targets(config, snapshot, target_keys=target_keys)
+
+
+def instantiate(
+    config: Any,
+    *args: Any,
+    mode: InstantiationMode = InstantiationMode.LENIENT,
+    **kwargs: Any,
+) -> Any:
+    """Instantiate a config after validating all targets against a stable allowlist."""
+    snapshot = _snapshot_allowlist()
+    _preflight_targets(config, snapshot)
+    _preflight_targets(kwargs, snapshot)
+    return _mlm_instantiate(config, *args, mode=mode, **kwargs)
+
+
+def instantiate_node(
+    node: Any,
+    *args: Any,
+    partial: bool = False,
+    mode: InstantiationMode = InstantiationMode.LENIENT,
+) -> Any:
+    """Instantiate an OmegaConf node after validating all targets against a stable allowlist."""
+    snapshot = _snapshot_allowlist()
+    _preflight_targets(node, snapshot)
+    return _mlm_instantiate_node(node, *args, partial=partial, mode=mode)
+
+
 def register_allowed_target_prefix(prefix: str) -> None:
     """Register an additional allowed module prefix for _target_ instantiation.
 

tests/unit_tests/training/test_model_load_save.py

@@ -28,12 +28,14 @@
     dtype_from_hf,
     dtype_from_str,
     load_megatron_model,
+    load_model_config,
     load_tokenizer,
     megatron_cpu_init_context,
     save_megatron_model,
     temporary_distributed_context,
     torch_dtype_from_mcore_config,
 )
+from megatron.bridge.utils.instantiate_utils import InstantiationException
 
 
 class TestTorchDtypeFromMcoreConfig:
@@ -270,7 +272,10 @@ def test_load_mbridge_saved_model_config(
         mock_dist.is_initialized.return_value = False
 
         mock_run_cfg_dict = {
-            "model": {"tensor_model_parallel_size": 1, "_builder_": "import.path.to.SomeModelBuilder"}
+            "model": {
+                "tensor_model_parallel_size": 1,
+                "_builder_": "megatron.bridge.models.gpt.gpt_builder.GPTModelBuilder",
+            }
         }
         mock_run_config.return_value = mock_run_cfg_dict
 
@@ -323,6 +328,63 @@ def test_load_mbridge_saved_model_config(
         assert result == [mock_model]
         mock_load_weights.assert_called_with(ckpt_path, [mock_model], return_state_dict=False)
 
+    @patch("megatron.bridge.training.checkpointing.read_run_config")
+    @patch("megatron.bridge.training.checkpointing.get_checkpoint_run_config_filename")
+    @patch("megatron.bridge.training.model_load_save.ModelConfig.from_dict")
+    def test_load_model_config_rejects_untrusted_nested_builder_metadata_target(
+        self,
+        mock_from_dict,
+        mock_run_config_fname,
+        mock_run_config,
+    ):
+        """Test that _builder_ checkpoint metadata cannot import untrusted nested targets."""
+        mock_run_config.return_value = {
+            "model": {
+                "_target_": "megatron.bridge.models.gpt.gpt_builder.GPTModelConfig",
+                "_builder_": "megatron.bridge.models.gpt.gpt_builder.GPTModelBuilder",
+                "transformer_layer_spec": {
+                    "_target_": "llama3_ckpt_interactive.iter_0000001.payload.ShellConfig",
+                },
+            }
+        }
+
+        with tempfile.TemporaryDirectory() as ckpt_path:
+            config_file = Path(ckpt_path) / "run_config.yaml"
+            config_file.touch()
+            mock_run_config_fname.return_value = str(config_file)
+
+            with pytest.raises(InstantiationException, match="not in the allowlist"):
+                load_model_config(ckpt_path)
+
+        mock_from_dict.assert_not_called()
+
+    @patch("megatron.bridge.training.checkpointing.read_run_config")
+    @patch("megatron.bridge.training.checkpointing.get_checkpoint_run_config_filename")
+    @patch("megatron.bridge.training.model_load_save.ModelConfig.from_dict")
+    def test_load_model_config_rejects_untrusted_builder_metadata_builder(
+        self,
+        mock_from_dict,
+        mock_run_config_fname,
+        mock_run_config,
+    ):
+        """Test that _builder_ checkpoint metadata cannot select an untrusted builder class."""
+        mock_run_config.return_value = {
+            "model": {
+                "_target_": "megatron.bridge.models.gpt.gpt_builder.GPTModelConfig",
+                "_builder_": "llama3_ckpt_interactive.iter_0000001.payload.ShellBuilder",
+            }
+        }
+
+        with tempfile.TemporaryDirectory() as ckpt_path:
+            config_file = Path(ckpt_path) / "run_config.yaml"
+            config_file.touch()
+            mock_run_config_fname.return_value = str(config_file)
+
+            with pytest.raises(InstantiationException, match="not in the allowlist"):
+                load_model_config(ckpt_path)
+
+        mock_from_dict.assert_not_called()
+
     @pytest.mark.parametrize("model_type", ["gpt", "mamba", "resnet"])
     @patch("megatron.bridge.training.model_load_save.temporary_distributed_context")
     @patch("megatron.bridge.training.mlm_compat.model._mamba_provider")