fix(lora): reject incomplete native adapter checkpoints

miles/backends/megatron_utils/lora_utils.py

@@ -521,6 +521,51 @@ def save_lora_checkpoint(
     return str(save_path)
 
 
+def _validate_native_adapter_state(
+    model: Sequence[torch.nn.Module],
+    state_dict: dict[str, torch.Tensor],
+    native_path: Path,
+) -> dict[str, torch.nn.Parameter]:
+    expected = {
+        name: param
+        for model_chunk in model
+        for name, param in model_chunk.named_parameters()
+        if _is_adapter_param_name(name)
+    }
+    if not expected:
+        # Otherwise an empty shard matches an empty model and the load reports success.
+        raise RuntimeError(f"Cannot load {native_path.name}: the model exposes no LoRA adapter parameters")
+
+    mismatched: list[str] = []
+    non_finite: list[str] = []
+    for name in expected.keys() & state_dict.keys():
+        tensor = state_dict[name]
+        if (
+            not isinstance(tensor, torch.Tensor)
+            or tensor.shape != expected[name].shape
+            or tensor.dtype != expected[name].dtype
+        ):
+            mismatched.append(name)
+        elif (tensor.is_floating_point() or tensor.is_complex()) and not torch.isfinite(tensor).all():
+            non_finite.append(name)
+
+    problems = [
+        f"{label} {len(names)}"
+        for label, names in (
+            ("missing", expected.keys() - state_dict.keys()),
+            ("unexpected", state_dict.keys() - expected.keys()),
+            ("dtype/shape mismatch", mismatched),
+            ("non-finite", non_finite),
+        )
+        if names
+    ]
+    if problems:
+        raise RuntimeError(
+            f"Native LoRA checkpoint {native_path.name} does not match the model: {', '.join(problems)}"
+        )
+    return expected
+
+
 def load_lora_adapter(
     model: Sequence[torch.nn.Module],
     adapter_path: str,
@@ -567,13 +612,10 @@ def load_lora_adapter(
             native_path = legacy
     if native_path.exists():
         state_dict = torch.load(native_path, map_location="cpu", weights_only=True)
-        loaded = 0
-        for model_chunk in model:
-            for name, param in model_chunk.named_parameters():
-                if name in state_dict:
-                    param.data.copy_(state_dict[name].to(device=param.device))
-                    loaded += 1
-        logger.info(f"Loaded {loaded} adapter tensors from Megatron-native checkpoint: {native_path}")
+        adapter_params = _validate_native_adapter_state(model, state_dict, native_path)
+        for name, param in adapter_params.items():
+            param.data.copy_(state_dict[name].to(device=param.device))
+        logger.info(f"Loaded {len(adapter_params)} adapter tensors from Megatron-native checkpoint: {native_path}")
 
         iteration = _load_training_state(adapter_dir, optimizer, opt_param_scheduler)
         return True, iteration

tests/fast/backends/megatron_utils/test_lora_utils.py

@@ -8,7 +8,9 @@
 from unittest.mock import MagicMock
 
 import pytest
+import torch
 
+import miles.backends.megatron_utils.lora_utils as lora_utils
 from miles.backends.megatron_utils.lora_utils import (
     _get_lora_class_name,
     _is_adapter_param_name,
@@ -348,6 +350,87 @@ def test_canonical_target_modules(self):
         assert config["r"] == 8
 
 
+# ---------------------------------------------------------------------------
+# Native adapter checkpoint validation
+# ---------------------------------------------------------------------------
+
+
+class TestNativeAdapterCheckpointValidation:
+    adapter_name = "module.decoder.layers.0.self_attention.linear_qkv.lora_A.weight"
+
+    @classmethod
+    def _model(cls):
+        param = torch.nn.Parameter(torch.ones(2, 2))
+        model = MagicMock()
+        model.named_parameters.return_value = [(cls.adapter_name, param)]
+        return model, param
+
+    @staticmethod
+    def _patch_parallel_state(monkeypatch):
+        parallel_state = MagicMock()
+        parallel_state.tp.rank = 0
+        parallel_state.pp.rank = 0
+        monkeypatch.setattr(lora_utils, "get_parallel_state", lambda: parallel_state)
+
+    def test_exact_native_shard_loads(self, tmp_path, monkeypatch):
+        self._patch_parallel_state(monkeypatch)
+        model, param = self._model()
+        expected = torch.full((2, 2), 3.0)
+        torch.save({self.adapter_name: expected}, tmp_path / "adapter_megatron_rank0.pt")
+
+        loaded, iteration = lora_utils.load_lora_adapter([model], str(tmp_path))
+
+        assert loaded
+        assert iteration is None
+        assert torch.equal(param, expected)
+
+    @pytest.mark.parametrize(
+        ("failure", "problem"),
+        [
+            ("missing", "missing 1"),
+            ("unexpected", "unexpected 1"),
+            ("shape", "dtype/shape mismatch 1"),
+            ("dtype", "dtype/shape mismatch 1"),
+            ("non_finite", "non-finite 1"),
+        ],
+    )
+    def test_invalid_native_shard_stops_before_weights_or_training_state(
+        self, failure, problem, tmp_path, monkeypatch
+    ):
+        self._patch_parallel_state(monkeypatch)
+        model, param = self._model()
+        state_dict = {self.adapter_name: torch.full((2, 2), 3.0)}
+        if failure == "missing":
+            state_dict.clear()
+        elif failure == "unexpected":
+            state_dict["unexpected.lora_A.weight"] = torch.ones(1)
+        elif failure == "shape":
+            state_dict[self.adapter_name] = torch.ones(3, 2)
+        elif failure == "dtype":
+            state_dict[self.adapter_name] = torch.ones(2, 2, dtype=torch.float16)
+        else:
+            state_dict[self.adapter_name] = torch.full((2, 2), float("nan"))
+        torch.save(state_dict, tmp_path / "adapter_megatron_rank0.pt")
+        torch.save({"optimizer": {}, "iteration": 7}, tmp_path / "training_state_rank0.pt")
+        optimizer = MagicMock()
+
+        with pytest.raises(RuntimeError, match=problem):
+            lora_utils.load_lora_adapter([model], str(tmp_path), optimizer=optimizer)
+
+        assert torch.equal(param, torch.ones(2, 2))
+        optimizer.load_state_dict.assert_not_called()
+
+    def test_model_without_adapter_parameters_is_not_a_successful_load(self, tmp_path, monkeypatch):
+        # An empty shard matches an empty model on every other check.
+        self._patch_parallel_state(monkeypatch)
+        model = MagicMock()
+        model.named_parameters.return_value = []
+        torch.save({}, tmp_path / "adapter_megatron_rank0.pt")
+
+        with pytest.raises(RuntimeError, match="exposes no LoRA adapter parameters"):
+            lora_utils.load_lora_adapter([model], str(tmp_path))
+
+
 # ---------------------------------------------------------------------------
 # LORA_ADAPTER_NAME constant
 # ---------------------------------------------------------------------------