[diffusion] refactor: resolve lora weight sources deterministically (#35774)
This commit is contained in:
Mick
@@ -52,6 +52,7 @@ from sglang.multimodal_gen.runtime.utils.model_overlay import (
|
||||
from sglang.multimodal_gen.runtime.utils.quantization_utils import (
|
||||
normalize_flat_modelopt_quant_config,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.weights.source import resolve_weight
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.utils.hf_transformers import check_gguf_file
|
||||
from sglang.utils import is_in_ci
|
||||
@@ -616,42 +617,66 @@ def maybe_download_lora(
|
||||
Returns:
|
||||
Local path to the model
|
||||
"""
|
||||
# Repositories often publish several adapter revisions side by side. If a
|
||||
# filename is pinned, do not download every weight before selecting it.
|
||||
# Keep JSON metadata so PEFT's lora_alpha remains available.
|
||||
allow_patterns = (
|
||||
["*.json", weight_name, f"**/{weight_name}"]
|
||||
if weight_name is not None
|
||||
else ["*.json", "*.safetensors", "*.bin"]
|
||||
)
|
||||
if envs.SGLANG_USE_MODELSCOPE.get():
|
||||
allow_patterns = (
|
||||
["*.json", weight_name, f"**/{weight_name}"]
|
||||
if weight_name is not None
|
||||
else ["*.json", "*.safetensors", "*.bin"]
|
||||
)
|
||||
local_path = maybe_download_model(
|
||||
model_name_or_path,
|
||||
local_dir,
|
||||
download,
|
||||
is_lora=True,
|
||||
allow_patterns=allow_patterns,
|
||||
)
|
||||
if os.path.isfile(local_path):
|
||||
return local_path
|
||||
if weight_name is not None:
|
||||
target = os.path.join(local_path, weight_name)
|
||||
if not os.path.isfile(target):
|
||||
raise FileNotFoundError(
|
||||
f"Specified lora_weight_name '{weight_name}' not found in "
|
||||
f"{local_path}"
|
||||
)
|
||||
return target
|
||||
guessed = _best_guess_weight_name(local_path, file_extension=".safetensors")
|
||||
if guessed is None and current_platform.is_rocm():
|
||||
guessed = _best_guess_weight_name(
|
||||
model_name_or_path, file_extension=".safetensors"
|
||||
)
|
||||
return os.path.join(local_path, guessed)
|
||||
|
||||
resolved_weight = resolve_weight(model_name_or_path, weight_name=weight_name)
|
||||
selected_file = resolved_weight.selected_file
|
||||
if not selected_file.endswith(".safetensors"):
|
||||
raise ValueError(
|
||||
"Native diffusion LoRA loading requires a safetensors file, got "
|
||||
f"{selected_file!r}"
|
||||
)
|
||||
|
||||
source = resolved_weight.inventory.source
|
||||
if source.kind == "local":
|
||||
assert source.local_path is not None
|
||||
if os.path.isfile(source.local_path):
|
||||
return source.local_path
|
||||
return os.path.join(source.local_path, selected_file)
|
||||
|
||||
assert source.repo_id is not None
|
||||
local_path = maybe_download_model(
|
||||
model_name_or_path,
|
||||
source.repo_id,
|
||||
local_dir,
|
||||
download,
|
||||
is_lora=True,
|
||||
allow_patterns=allow_patterns,
|
||||
allow_patterns=["*.json", selected_file],
|
||||
revision=resolved_weight.inventory.resolved_revision or source.revision,
|
||||
)
|
||||
# return directly if local_path is a file
|
||||
if os.path.isfile(local_path):
|
||||
return local_path
|
||||
|
||||
if weight_name is not None:
|
||||
target = os.path.join(local_path, weight_name)
|
||||
if not os.path.isfile(target):
|
||||
raise FileNotFoundError(
|
||||
f"Specified lora_weight_name '{weight_name}' not found in {local_path}"
|
||||
)
|
||||
return target
|
||||
|
||||
guessed = _best_guess_weight_name(local_path, file_extension=".safetensors")
|
||||
# AMD workaround: PR 15813 changed from model_name_or_path to local_path,
|
||||
# which can return None. Fall back to original behavior on ROCm.
|
||||
if guessed is None and current_platform.is_rocm():
|
||||
guessed = _best_guess_weight_name(
|
||||
model_name_or_path, file_extension=".safetensors"
|
||||
target = os.path.join(local_path, selected_file)
|
||||
if not os.path.isfile(target):
|
||||
raise FileNotFoundError(
|
||||
f"Resolved LoRA weight {selected_file!r} was not downloaded to {local_path}"
|
||||
)
|
||||
return os.path.join(local_path, guessed)
|
||||
return target
|
||||
|
||||
|
||||
def verify_model_config_and_directory(model_path: str) -> dict[str, Any]:
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
"""Weight source contracts shared by runtime loaders."""
|
||||
@@ -0,0 +1,266 @@
|
||||
"""Resolve weight sources for runtime loaders."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path, PurePosixPath
|
||||
from typing import Literal
|
||||
from urllib.parse import unquote, urlparse
|
||||
|
||||
from huggingface_hub import HfApi
|
||||
from huggingface_hub.utils import validate_repo_id
|
||||
|
||||
WeightSourceKind = Literal["local", "huggingface"]
|
||||
|
||||
_WEIGHT_SUFFIXES = (".safetensors", ".gguf", ".bin", ".pt", ".pth", ".ckpt")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class WeightSource:
|
||||
original: str
|
||||
kind: WeightSourceKind
|
||||
local_path: str | None = None
|
||||
repo_id: str | None = None
|
||||
revision: str | None = None
|
||||
subfolder: str | None = None
|
||||
filename: str | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class WeightInventory:
|
||||
source: WeightSource
|
||||
resolved_revision: str | None
|
||||
files: tuple[str, ...]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ResolvedWeight:
|
||||
inventory: WeightInventory
|
||||
selected_file: str
|
||||
|
||||
|
||||
def _validate_relative_hub_path(path: str, field_name: str) -> str:
|
||||
normalized = str(PurePosixPath(path))
|
||||
pure_path = PurePosixPath(normalized)
|
||||
if not path or pure_path.is_absolute() or ".." in pure_path.parts:
|
||||
raise ValueError(f"Invalid Hugging Face {field_name}: {path!r}")
|
||||
return normalized
|
||||
|
||||
|
||||
def _merge_revision(url_revision: str | None, revision: str | None) -> str | None:
|
||||
if url_revision is not None and revision is not None and url_revision != revision:
|
||||
raise ValueError(
|
||||
f"Weight URL pins revision {url_revision!r}, which conflicts with "
|
||||
f"revision {revision!r}"
|
||||
)
|
||||
return url_revision or revision
|
||||
|
||||
|
||||
def _parse_huggingface_url(source: str, revision: str | None) -> WeightSource:
|
||||
parsed = urlparse(source)
|
||||
if parsed.netloc.lower() not in ("huggingface.co", "www.huggingface.co"):
|
||||
raise ValueError(
|
||||
"Only huggingface.co weight URLs are supported; use a local path "
|
||||
"or an owner/repo reference for other sources"
|
||||
)
|
||||
|
||||
raw_parts = [part for part in parsed.path.split("/") if part]
|
||||
if raw_parts and raw_parts[0] in ("datasets", "spaces"):
|
||||
raise ValueError("Diffusion weights must come from a Hugging Face model repo")
|
||||
if len(raw_parts) < 2:
|
||||
raise ValueError(f"Hugging Face weight URL has no model repo: {source!r}")
|
||||
|
||||
repo_id = "/".join(unquote(part) for part in raw_parts[:2])
|
||||
validate_repo_id(repo_id)
|
||||
action = raw_parts[2] if len(raw_parts) > 2 else None
|
||||
if action is None:
|
||||
return WeightSource(
|
||||
original=source,
|
||||
kind="huggingface",
|
||||
repo_id=repo_id,
|
||||
revision=revision,
|
||||
)
|
||||
if action not in ("tree", "blob", "resolve") or len(raw_parts) < 4:
|
||||
raise ValueError(f"Unsupported Hugging Face weight URL: {source!r}")
|
||||
|
||||
url_revision = unquote(raw_parts[3])
|
||||
selected_revision = _merge_revision(url_revision, revision)
|
||||
tail = "/".join(unquote(part) for part in raw_parts[4:])
|
||||
if action == "tree":
|
||||
subfolder = _validate_relative_hub_path(tail, "subfolder") if tail else None
|
||||
return WeightSource(
|
||||
original=source,
|
||||
kind="huggingface",
|
||||
repo_id=repo_id,
|
||||
revision=selected_revision,
|
||||
subfolder=subfolder,
|
||||
)
|
||||
if not tail:
|
||||
raise ValueError(f"Hugging Face weight URL has no filename: {source!r}")
|
||||
return WeightSource(
|
||||
original=source,
|
||||
kind="huggingface",
|
||||
repo_id=repo_id,
|
||||
revision=selected_revision,
|
||||
filename=_validate_relative_hub_path(tail, "filename"),
|
||||
)
|
||||
|
||||
|
||||
def parse_weight_source(
|
||||
source: str,
|
||||
*,
|
||||
revision: str | None = None,
|
||||
) -> WeightSource:
|
||||
"""Parse local paths, Hub repo IDs, subfolders, and exact Hub URLs."""
|
||||
expanded = os.path.expanduser(source)
|
||||
parsed = urlparse(source)
|
||||
if parsed.scheme in ("http", "https"):
|
||||
return _parse_huggingface_url(source, revision)
|
||||
|
||||
looks_local = (
|
||||
os.path.exists(expanded)
|
||||
or os.path.isabs(expanded)
|
||||
or source.startswith(("./", "../", "~"))
|
||||
)
|
||||
if looks_local:
|
||||
return WeightSource(
|
||||
original=source,
|
||||
kind="local",
|
||||
local_path=os.path.abspath(expanded),
|
||||
)
|
||||
|
||||
parts = source.split("/")
|
||||
if len(parts) < 2 or not all(parts[:2]):
|
||||
raise ValueError(
|
||||
f"Weight source {source!r} is neither a local path nor an "
|
||||
"owner/repo Hugging Face reference"
|
||||
)
|
||||
repo_id = "/".join(parts[:2])
|
||||
validate_repo_id(repo_id)
|
||||
tail = "/".join(parts[2:]) or None
|
||||
filename = (
|
||||
_validate_relative_hub_path(tail, "filename")
|
||||
if tail is not None and tail.lower().endswith(_WEIGHT_SUFFIXES)
|
||||
else None
|
||||
)
|
||||
subfolder = tail if filename is None else None
|
||||
if subfolder is not None:
|
||||
subfolder = _validate_relative_hub_path(subfolder, "subfolder")
|
||||
return WeightSource(
|
||||
original=source,
|
||||
kind="huggingface",
|
||||
repo_id=repo_id,
|
||||
revision=revision,
|
||||
subfolder=subfolder,
|
||||
filename=filename,
|
||||
)
|
||||
|
||||
|
||||
def _filter_inventory_files(
|
||||
files: tuple[str, ...], source: WeightSource
|
||||
) -> tuple[str, ...]:
|
||||
if source.filename is not None:
|
||||
selected = tuple(path for path in files if path == source.filename)
|
||||
if not selected:
|
||||
raise FileNotFoundError(
|
||||
f"Weight file {source.filename!r} was not found in {source.repo_id}"
|
||||
)
|
||||
return selected
|
||||
if source.subfolder is None:
|
||||
return files
|
||||
prefix = source.subfolder.rstrip("/") + "/"
|
||||
selected = tuple(path for path in files if path.startswith(prefix))
|
||||
if not selected:
|
||||
raise FileNotFoundError(
|
||||
f"Weight subfolder {source.subfolder!r} was not found in {source.repo_id}"
|
||||
)
|
||||
return selected
|
||||
|
||||
|
||||
def resolve_weight_inventory(source: WeightSource) -> WeightInventory:
|
||||
"""List source files and pin a remote source to an immutable revision."""
|
||||
if source.kind == "local":
|
||||
assert source.local_path is not None
|
||||
local_path = Path(source.local_path)
|
||||
if not local_path.exists():
|
||||
raise FileNotFoundError(f"Weight path does not exist: {local_path}")
|
||||
if local_path.is_file():
|
||||
files = (local_path.name,)
|
||||
else:
|
||||
files = tuple(
|
||||
path.relative_to(local_path).as_posix()
|
||||
for path in sorted(local_path.rglob("*"))
|
||||
if path.is_file()
|
||||
)
|
||||
return WeightInventory(
|
||||
source=source,
|
||||
resolved_revision=None,
|
||||
files=files,
|
||||
)
|
||||
|
||||
assert source.repo_id is not None
|
||||
model_info = HfApi().model_info(
|
||||
source.repo_id,
|
||||
revision=source.revision,
|
||||
)
|
||||
files = tuple(sibling.rfilename for sibling in model_info.siblings)
|
||||
return WeightInventory(
|
||||
source=source,
|
||||
resolved_revision=model_info.sha,
|
||||
files=_filter_inventory_files(files, source),
|
||||
)
|
||||
|
||||
|
||||
def _select_named_file(candidates: tuple[str, ...], weight_name: str) -> str:
|
||||
exact = tuple(path for path in candidates if path == weight_name)
|
||||
if exact:
|
||||
return exact[0]
|
||||
basename_matches = tuple(
|
||||
path for path in candidates if PurePosixPath(path).name == weight_name
|
||||
)
|
||||
if len(basename_matches) == 1:
|
||||
return basename_matches[0]
|
||||
if not basename_matches:
|
||||
raise FileNotFoundError(f"Requested weight {weight_name!r} was not found")
|
||||
raise ValueError(
|
||||
f"Weight name {weight_name!r} matches multiple files: "
|
||||
f"{list(basename_matches)}"
|
||||
)
|
||||
|
||||
|
||||
def select_weight_file(
|
||||
inventory: WeightInventory, weight_name: str | None = None
|
||||
) -> str:
|
||||
"""Select weights deterministically; never guess among independent files."""
|
||||
candidates = tuple(
|
||||
path for path in inventory.files if path.lower().endswith(_WEIGHT_SUFFIXES)
|
||||
)
|
||||
if inventory.source.filename is not None:
|
||||
return inventory.files[0]
|
||||
if weight_name is not None:
|
||||
return _select_named_file(candidates, weight_name)
|
||||
if len(candidates) == 1:
|
||||
return candidates[0]
|
||||
if not candidates:
|
||||
raise FileNotFoundError("Source contains no recognized weight files")
|
||||
raise ValueError(
|
||||
"Source contains multiple independent weight files; select one with "
|
||||
f"an exact file URL or weight name. Candidates: {list(candidates)}"
|
||||
)
|
||||
|
||||
|
||||
def resolve_weight(
|
||||
source: str,
|
||||
*,
|
||||
revision: str | None = None,
|
||||
weight_name: str | None = None,
|
||||
) -> ResolvedWeight:
|
||||
"""Resolve one weight file without downloading its tensor payload."""
|
||||
parsed_source = parse_weight_source(source, revision=revision)
|
||||
inventory = resolve_weight_inventory(parsed_source)
|
||||
selected_file = select_weight_file(inventory, weight_name)
|
||||
return ResolvedWeight(
|
||||
inventory=inventory,
|
||||
selected_file=selected_file,
|
||||
)
|
||||
@@ -142,20 +142,71 @@ def test_lora_alpha_override_updates_cached_adapter_scale():
|
||||
assert layer.lora_alpha == 8
|
||||
|
||||
|
||||
def test_pinned_lora_weight_limits_snapshot_download(tmp_path):
|
||||
def test_lora_tree_url_selects_one_pinned_weight(tmp_path):
|
||||
weight_name = "adapter-v4.safetensors"
|
||||
weight_path = tmp_path / weight_name
|
||||
adapter_dir = tmp_path / "adapters"
|
||||
adapter_dir.mkdir()
|
||||
weight_path = adapter_dir / weight_name
|
||||
weight_path.touch()
|
||||
model_info = SimpleNamespace(
|
||||
sha="immutable-sha",
|
||||
siblings=[
|
||||
SimpleNamespace(rfilename="adapters/adapter-v3.safetensors"),
|
||||
SimpleNamespace(rfilename="adapters/adapter-v4.safetensors"),
|
||||
],
|
||||
)
|
||||
|
||||
download_target = (
|
||||
"sglang.multimodal_gen.runtime.utils.hf_diffusers_utils.maybe_download_model"
|
||||
)
|
||||
with patch(download_target, return_value=str(tmp_path)) as download:
|
||||
actual = maybe_download_lora("org/multi-adapter", weight_name=weight_name)
|
||||
with (
|
||||
patch(
|
||||
"sglang.multimodal_gen.runtime.weights.source.HfApi.model_info",
|
||||
return_value=model_info,
|
||||
),
|
||||
patch(download_target, return_value=str(tmp_path)) as download,
|
||||
):
|
||||
actual = maybe_download_lora(
|
||||
"https://huggingface.co/org/multi-adapter/tree/main/adapters",
|
||||
weight_name=weight_name,
|
||||
)
|
||||
|
||||
assert actual == str(weight_path)
|
||||
assert download.call_args.args[0] == "org/multi-adapter"
|
||||
assert download.call_args.kwargs["revision"] == "immutable-sha"
|
||||
assert download.call_args.kwargs["allow_patterns"] == [
|
||||
"*.json",
|
||||
f"adapters/{weight_name}",
|
||||
]
|
||||
|
||||
|
||||
def test_lora_exact_file_url_needs_no_weight_name(tmp_path):
|
||||
weight_path = tmp_path / "adapter.safetensors"
|
||||
weight_path.touch()
|
||||
model_info = SimpleNamespace(
|
||||
sha="immutable-sha",
|
||||
siblings=[
|
||||
SimpleNamespace(rfilename="adapter.safetensors"),
|
||||
SimpleNamespace(rfilename="other.safetensors"),
|
||||
],
|
||||
)
|
||||
|
||||
download_target = (
|
||||
"sglang.multimodal_gen.runtime.utils.hf_diffusers_utils.maybe_download_model"
|
||||
)
|
||||
with (
|
||||
patch(
|
||||
"sglang.multimodal_gen.runtime.weights.source.HfApi.model_info",
|
||||
return_value=model_info,
|
||||
),
|
||||
patch(download_target, return_value=str(tmp_path)) as download,
|
||||
):
|
||||
actual = maybe_download_lora(
|
||||
"https://huggingface.co/org/multi-adapter/resolve/main/adapter.safetensors"
|
||||
)
|
||||
|
||||
assert actual == str(weight_path)
|
||||
assert download.call_args.kwargs["allow_patterns"] == [
|
||||
"*.json",
|
||||
weight_name,
|
||||
f"**/{weight_name}",
|
||||
"adapter.safetensors",
|
||||
]
|
||||
|
||||
@@ -0,0 +1,82 @@
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from sglang.multimodal_gen.runtime.weights.source import (
|
||||
parse_weight_source,
|
||||
resolve_weight,
|
||||
resolve_weight_inventory,
|
||||
)
|
||||
|
||||
|
||||
def test_parse_weight_source_accepts_repo_subfolder_and_exact_url():
|
||||
subfolder = parse_weight_source("owner/repo/text_encoder", revision="v1")
|
||||
repo_file = parse_weight_source("owner/repo/adapter.safetensors")
|
||||
exact_file = parse_weight_source(
|
||||
"https://huggingface.co/owner/repo/resolve/main/weights/model.safetensors"
|
||||
)
|
||||
|
||||
assert subfolder.repo_id == "owner/repo"
|
||||
assert subfolder.subfolder == "text_encoder"
|
||||
assert subfolder.revision == "v1"
|
||||
assert repo_file.filename == "adapter.safetensors"
|
||||
assert repo_file.subfolder is None
|
||||
assert exact_file.repo_id == "owner/repo"
|
||||
assert exact_file.revision == "main"
|
||||
assert exact_file.filename == "weights/model.safetensors"
|
||||
|
||||
|
||||
def test_parse_weight_source_rejects_conflicting_url_revision():
|
||||
with pytest.raises(ValueError, match="conflicts with revision"):
|
||||
parse_weight_source(
|
||||
"https://huggingface.co/owner/repo/tree/main/transformer",
|
||||
revision="v2",
|
||||
)
|
||||
|
||||
|
||||
def test_resolve_local_inventory_lists_files_without_loading_tensors(tmp_path):
|
||||
component = tmp_path / "component"
|
||||
component.mkdir()
|
||||
(component / "config.json").write_text("{}")
|
||||
(component / "model.safetensors").write_bytes(b"header-only-fixture")
|
||||
|
||||
inventory = resolve_weight_inventory(parse_weight_source(str(component)))
|
||||
|
||||
assert inventory.resolved_revision is None
|
||||
assert list(inventory.files) == [
|
||||
"config.json",
|
||||
"model.safetensors",
|
||||
]
|
||||
|
||||
|
||||
def test_resolve_remote_inventory_pins_revision_and_filters_subfolder():
|
||||
source = parse_weight_source("owner/repo/text_encoder", revision="main")
|
||||
model_info = SimpleNamespace(
|
||||
sha="immutable-sha",
|
||||
siblings=[
|
||||
SimpleNamespace(rfilename="text_encoder/config.json"),
|
||||
SimpleNamespace(rfilename="text_encoder/model.safetensors"),
|
||||
SimpleNamespace(rfilename="vae/config.json"),
|
||||
],
|
||||
)
|
||||
|
||||
with patch(
|
||||
"sglang.multimodal_gen.runtime.weights.source.HfApi.model_info",
|
||||
return_value=model_info,
|
||||
):
|
||||
inventory = resolve_weight_inventory(source)
|
||||
|
||||
assert inventory.resolved_revision == "immutable-sha"
|
||||
assert inventory.files == (
|
||||
"text_encoder/config.json",
|
||||
"text_encoder/model.safetensors",
|
||||
)
|
||||
|
||||
|
||||
def test_weight_source_rejects_ambiguous_files(tmp_path):
|
||||
(tmp_path / "a.safetensors").write_bytes(b"a")
|
||||
(tmp_path / "b.safetensors").write_bytes(b"b")
|
||||
|
||||
with pytest.raises(ValueError, match="multiple independent weight files"):
|
||||
resolve_weight(str(tmp_path))
|
||||
Reference in New Issue
Block a user