diff --git a/miles/backends/megatron_utils/lora_utils.py b/miles/backends/megatron_utils/lora_utils.py index b86eb518ae4..e59923819d7 100644 --- a/miles/backends/megatron_utils/lora_utils.py +++ b/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 diff --git a/tests/fast/backends/megatron_utils/test_lora_utils.py b/tests/fast/backends/megatron_utils/test_lora_utils.py index ca902420b40..3b671034f1d 100644 --- a/tests/fast/backends/megatron_utils/test_lora_utils.py +++ b/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 # ---------------------------------------------------------------------------