diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/lora/peft_adapter.py b/python/sglang/multimodal_gen/runtime/pipelines_core/lora/peft_adapter.py index 840e37174..a6aed9cfa 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/lora/peft_adapter.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/lora/peft_adapter.py @@ -9,6 +9,7 @@ from pathlib import Path from typing import Any, Mapping import torch +from safetensors import safe_open _ADAPTER_SLOT = re.compile(r"(\.lora_[AB])\.([^.]+)\.weight$") _WRAPPER_PREFIXES = ("peft_model.base_model.model.", "base_model.model.") @@ -21,16 +22,78 @@ _UNSUPPORTED_CONFIG_FIELDS = ( "use_bdlora", "use_qalora", ) +_SAFETENSORS_ALPHA_KEYS = ("lora_alpha", "network_alpha", "alpha") +_NATIVE_LORA_A_SUFFIXES = ( + ".lora_A.weight", + ".lora_down.weight", + ".lora.down.weight", +) + + +def _has_unambiguous_global_alpha(file: Any) -> bool: + """Reject bare mixed-rank files whose global alpha semantics are ambiguous.""" + keys = list(file.keys()) + if any(_ADAPTER_SLOT.search(name) is not None for name in keys): + return True + + ranks = set() + for name in keys: + if not name.endswith(_NATIVE_LORA_A_SUFFIXES): + continue + shape = file.get_slice(name).get_shape() + if len(shape) >= 2: + ranks.add(shape[-2]) + return len(ranks) == 1 + + +def _load_safetensors_lora_alpha(weight_path: str) -> int | None: + if Path(weight_path).suffix.lower() != ".safetensors": + return None + with safe_open(weight_path, framework="pt", device="cpu") as file: + metadata = file.metadata() or {} + if not _has_unambiguous_global_alpha(file): + return None + declared = [] + for key in _SAFETENSORS_ALPHA_KEYS: + value = metadata.get(key) + if value is None: + continue + try: + numeric = float(value) + except (TypeError, ValueError) as error: + raise ValueError( + f"safetensors metadata {key!r} must be a positive integer" + ) from error + if not math.isfinite(numeric) or numeric <= 0 or not numeric.is_integer(): + raise ValueError(f"safetensors metadata {key!r} must be a positive integer") + declared.append((key, int(numeric))) + values = {value for _, value in declared} + if len(values) > 1: + raise ValueError(f"conflicting safetensors LoRA alpha metadata: {declared}") + return declared[0][1] if declared else None def load_peft_config(weight_path: str) -> dict[str, Any]: path = Path(weight_path).with_name("adapter_config.json") - if not path.is_file(): - return {} - with path.open(encoding="utf-8") as file: - config = json.load(file) + config = {} + if path.is_file(): + with path.open(encoding="utf-8") as file: + config = json.load(file) if not isinstance(config, dict): raise ValueError("PEFT adapter_config.json must contain a JSON object") + metadata_alpha = _load_safetensors_lora_alpha(weight_path) + config_alpha = get_peft_lora_alpha(config) + if ( + metadata_alpha is not None + and config_alpha is not None + and metadata_alpha != config_alpha + ): + raise ValueError( + "adapter_config.json lora_alpha conflicts with safetensors metadata: " + f"{config_alpha} != {metadata_alpha}" + ) + if metadata_alpha is not None: + config.setdefault("lora_alpha", metadata_alpha) return config diff --git a/python/sglang/multimodal_gen/test/unit/test_lora_peft.py b/python/sglang/multimodal_gen/test/unit/test_lora_peft.py index 56e4a518a..4ce88d050 100644 --- a/python/sglang/multimodal_gen/test/unit/test_lora_peft.py +++ b/python/sglang/multimodal_gen/test/unit/test_lora_peft.py @@ -4,12 +4,14 @@ import math import pytest import torch +from safetensors.torch import save_file from sglang.multimodal_gen.runtime.pipelines_core.lora.format_adapter import ( normalize_lora_state_dict, ) from sglang.multimodal_gen.runtime.pipelines_core.lora.peft_adapter import ( get_peft_lora_alpha, + load_peft_config, ) @@ -53,3 +55,28 @@ def test_unsupported_peft_runtime_semantics_fail_closed(state_dict, adapter_conf def test_invalid_peft_lora_alpha_fails_closed(): with pytest.raises(ValueError, match="positive integer"): get_peft_lora_alpha({"lora_alpha": 8.5}) + + +def test_safetensors_alpha_metadata_supplies_peft_config(tmp_path): + weight_path = tmp_path / "adapter.safetensors" + save_file( + {"proj.lora_A.default.weight": torch.ones(4, 8)}, + weight_path, + metadata={"alpha": "128"}, + ) + + assert load_peft_config(str(weight_path))["lora_alpha"] == 128 + + +def test_mixed_rank_native_safetensors_does_not_apply_global_alpha(tmp_path): + weight_path = tmp_path / "adapter.safetensors" + save_file( + { + "gate.lora_A.weight": torch.ones(4, 8), + "proj.lora_A.weight": torch.ones(8, 8), + }, + weight_path, + metadata={"lora_alpha": "8"}, + ) + + assert "lora_alpha" not in load_peft_config(str(weight_path))