1490 lines
56 KiB
Python
1490 lines
56 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""
|
|
Tests for the DiffusersPipelineLoader.
|
|
"""
|
|
|
|
import json
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
import torch
|
|
import torch.nn as nn
|
|
from huggingface_hub import snapshot_download
|
|
from safetensors.torch import save_file
|
|
from vllm.config.load import LoadConfig
|
|
|
|
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_module
|
|
from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none
|
|
from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
|
|
from vllm_omni.diffusion.lora.manager import LoRABackend
|
|
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
|
from vllm_omni.diffusion.model_loader.host_weight_plan import (
|
|
HostWeightPlan,
|
|
HostWeightPlanResult,
|
|
TensorBinding,
|
|
)
|
|
from vllm_omni.diffusion.model_loader.host_weights import source_identity as source_identity_module
|
|
from vllm_omni.diffusion.models.helios import HeliosPipeline
|
|
from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract
|
|
from vllm_omni.diffusion.registry import initialize_model
|
|
from vllm_omni.quantization.component_config import ComponentQuantizationConfig
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
|
|
|
model_path = "hf-internal-testing/tiny-helios-modular-pipe"
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def prefetch_helios_model():
|
|
"""Downloads the tiny helios model prior to running a test."""
|
|
snapshot_download(model_path)
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
def mock_tp_group(mocker):
|
|
"""Mocks the tensor parallel group; this is needed to initialize the Helios model."""
|
|
mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", return_value=1)
|
|
mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_rank", return_value=0)
|
|
mock_group = mocker.MagicMock()
|
|
mock_group.world_size = 1
|
|
mock_group.rank_in_group = 0
|
|
mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_group)
|
|
|
|
|
|
class _DummyPipelineModel(nn.Module):
|
|
def __init__(self, *, source_prefix: str):
|
|
super().__init__()
|
|
self.transformer = nn.Linear(2, 2, bias=False)
|
|
self.vae = nn.Linear(2, 2, bias=False)
|
|
self.weights_sources = [
|
|
DiffusersPipelineLoader.ComponentSource(
|
|
model_or_path="dummy",
|
|
subfolder="transformer",
|
|
revision=None,
|
|
prefix=source_prefix,
|
|
fall_back_to_pt=True,
|
|
)
|
|
]
|
|
|
|
def load_weights(self, weights):
|
|
params = dict(self.named_parameters())
|
|
loaded: set[str] = set()
|
|
for name, tensor in weights:
|
|
if name not in params:
|
|
continue
|
|
params[name].data.copy_(tensor.to(dtype=params[name].dtype))
|
|
loaded.add(name)
|
|
return loaded
|
|
|
|
|
|
class _HWRTransformer(nn.Module):
|
|
host_weight_restore_contract = FinalLayoutModelContract(
|
|
implementation_id="test-hwr-transformer",
|
|
version="1",
|
|
)
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.weight = nn.Parameter(torch.arange(4, dtype=torch.float32).to(torch.bfloat16).reshape(2, 2))
|
|
|
|
def validate_restored_host_weights(self):
|
|
assert self.weight.dtype is torch.bfloat16
|
|
|
|
|
|
class _HWRPipeline(nn.Module):
|
|
def __init__(self, source_root: Path):
|
|
super().__init__()
|
|
self.transformer = _HWRTransformer()
|
|
self.load_count = 0
|
|
self.weights_sources = [
|
|
DiffusersPipelineLoader.ComponentSource(
|
|
model_or_path=str(source_root),
|
|
subfolder=None,
|
|
revision=None,
|
|
prefix="transformer.",
|
|
fall_back_to_pt=False,
|
|
)
|
|
]
|
|
|
|
def load_weights(self, weights):
|
|
loaded: set[str] = set()
|
|
params = dict(self.named_parameters())
|
|
for name, tensor in weights:
|
|
if name in params:
|
|
params[name].data.copy_(tensor.to(dtype=params[name].dtype))
|
|
loaded.add(name)
|
|
self.load_count += 1
|
|
return loaded
|
|
|
|
|
|
def _hwr_config(model: str | Path, root: Path, *, mode: str = "preferred") -> SimpleNamespace:
|
|
return SimpleNamespace(
|
|
model=str(model),
|
|
dtype=torch.bfloat16,
|
|
host_weight_runtime_mode=mode,
|
|
host_weight_runtime_root=str(root),
|
|
enable_distributed_layerwise_offload=True,
|
|
dlo_use_allgather=False,
|
|
lora_path=None,
|
|
quantization_config=None,
|
|
diffusion_attention_config=None,
|
|
parallel_config=SimpleNamespace(
|
|
use_hsdp=False,
|
|
data_parallel_size=1,
|
|
sequence_parallel_size=1,
|
|
tensor_parallel_size=1,
|
|
pipeline_parallel_size=1,
|
|
cfg_parallel_size=1,
|
|
enable_expert_parallel=False,
|
|
ulysses_degree=1,
|
|
ring_degree=1,
|
|
allgather_degree=1,
|
|
ulysses_mode="strict",
|
|
),
|
|
)
|
|
|
|
|
|
def _make_loader_with_weights(weight_names: list[str]) -> DiffusersPipelineLoader:
|
|
od_config = SimpleNamespace(
|
|
dtype=torch.float32,
|
|
parallel_config=SimpleNamespace(use_hsdp=False),
|
|
quantization_config=None,
|
|
)
|
|
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
|
|
|
|
loader.counter_before_loading_weights = 0.0
|
|
loader.counter_after_loading_weights = 0.0
|
|
|
|
def _iter_weights(_model):
|
|
for name in weight_names:
|
|
yield name, torch.zeros((2, 2))
|
|
|
|
loader.get_all_weights = _iter_weights # type: ignore[assignment]
|
|
return loader
|
|
|
|
|
|
def test_serialized_torchao_component_uses_pytorch_iterator(mocker):
|
|
torchao_config = SimpleNamespace(
|
|
get_name=lambda: "torchao",
|
|
is_checkpoint_torchao_serialized=True,
|
|
)
|
|
loader = _make_loader_with_weights([])
|
|
loader.quant_config = ComponentQuantizationConfig({"transformer": torchao_config})
|
|
loader.load_config.pt_load_map_location = "cpu"
|
|
files = ["/weights/model-10.bin", "/weights/model-2.bin"]
|
|
mocker.patch.object(loader, "_prepare_weights", return_value=("/weights", files, False))
|
|
pt_iterator = mocker.patch.object(
|
|
loader_module,
|
|
"pt_weights_iterator",
|
|
return_value=iter([("block.weight", torch.ones(1))]),
|
|
)
|
|
safetensors_iterator = mocker.patch.object(loader_module, "safetensors_weights_iterator")
|
|
multithread_iterator = mocker.patch.object(loader_module, "multi_thread_safetensors_weights_iterator")
|
|
source = DiffusersPipelineLoader.ComponentSource(
|
|
model_or_path="/weights",
|
|
subfolder=None,
|
|
revision=None,
|
|
prefix="transformer.",
|
|
)
|
|
|
|
weights = list(loader._get_weights_iterator(source))
|
|
|
|
pt_iterator.assert_called_once_with(
|
|
["/weights/model-2.bin", "/weights/model-10.bin"],
|
|
loader.load_config.use_tqdm_on_load,
|
|
"cpu",
|
|
)
|
|
safetensors_iterator.assert_not_called()
|
|
multithread_iterator.assert_not_called()
|
|
assert weights[0][0] == "transformer.block.weight"
|
|
assert torch.equal(weights[0][1], torch.ones(1))
|
|
|
|
|
|
def test_hwr_cold_publication_and_warm_restore_skip_ordinary_dit_loading(
|
|
tmp_path: Path,
|
|
monkeypatch,
|
|
):
|
|
canonical_root = tmp_path / "canonical"
|
|
canonical_root.mkdir()
|
|
save_file(
|
|
{"weight": torch.arange(4, dtype=torch.float32).to(torch.bfloat16).reshape(2, 2)},
|
|
str(canonical_root / "model.safetensors"),
|
|
)
|
|
store_root = tmp_path / "hwr-store"
|
|
hash_calls = 0
|
|
original_sha256 = source_identity_module._sha256_file
|
|
|
|
def counted_sha256(path: Path, state: object) -> str:
|
|
nonlocal hash_calls
|
|
hash_calls += 1
|
|
return original_sha256(path, state) # type: ignore[arg-type]
|
|
|
|
monkeypatch.setattr(source_identity_module, "_sha256_file", counted_sha256)
|
|
|
|
def make_loader() -> tuple[DiffusersPipelineLoader, _HWRPipeline]:
|
|
loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config(canonical_root, store_root))
|
|
pipeline = _HWRPipeline(canonical_root)
|
|
monkeypatch.setattr(loader, "_init_from_load_format", lambda *args, **kwargs: pipeline)
|
|
return loader, pipeline
|
|
|
|
cold_loader, cold_model = make_loader()
|
|
cold = cold_loader.load_model(load_device="cpu", device=torch.device("cpu"))
|
|
assert cold is cold_model
|
|
assert cold_model.load_count == 1
|
|
assert cold_loader._hwr_state is not None
|
|
|
|
warm_loader, warm_model = make_loader()
|
|
monkeypatch.setattr(
|
|
warm_loader,
|
|
"_process_weights_after_loading",
|
|
lambda *args, **kwargs: pytest.fail("warm HWR restore re-entered byte-changing finalization"),
|
|
)
|
|
warm = warm_loader.load_model(load_device="cpu", device=torch.device("cpu"))
|
|
|
|
assert warm is warm_model
|
|
assert warm_model.load_count == 0
|
|
assert torch.equal(warm_model.transformer.weight, cold_model.transformer.weight)
|
|
from vllm_omni.diffusion.offloader.startup import take_offload_startup_state
|
|
|
|
startup_state = take_offload_startup_state(warm)
|
|
assert startup_state is not None
|
|
warm_plan = startup_state.host_weight_plan
|
|
assert warm_plan is not None
|
|
assert warm_plan.lease_carrier is not None
|
|
warm_plan.lease_carrier.close()
|
|
assert hash_calls == 1
|
|
assert len(tuple((store_root / "source-digests-v1" / "entries").glob("*.json"))) == 1
|
|
|
|
|
|
def test_maybe_fuse_distilled_lora_skips_when_hwr_warm_snapshot_present():
|
|
cfg = SimpleNamespace(
|
|
lora_backend="distill",
|
|
lora_path="/path/to/lora.safetensors",
|
|
lora_scale=1.0,
|
|
dtype=torch.bfloat16,
|
|
quantization_config=None,
|
|
parallel_config=SimpleNamespace(use_hsdp=False),
|
|
)
|
|
loader = DiffusersPipelineLoader(LoadConfig(), cfg)
|
|
loader._hwr_state = {"warm_snapshot": {"transformer.weight": (1234, "hash")}}
|
|
|
|
model = nn.Module()
|
|
model.load_lora_weights = MagicMock()
|
|
|
|
loader._maybe_fuse_distilled_lora(model)
|
|
|
|
model.load_lora_weights.assert_not_called()
|
|
assert getattr(model, "lora_is_fused", False) is True
|
|
|
|
|
|
def test_maybe_fuse_distilled_lora_fuses_when_not_warm_snapshot():
|
|
cfg = SimpleNamespace(
|
|
lora_backend=LoRABackend.DISTILL,
|
|
lora_path="/path/to/lora.safetensors",
|
|
lora_scale=1.0,
|
|
dtype=torch.bfloat16,
|
|
quantization_config=None,
|
|
parallel_config=SimpleNamespace(use_hsdp=False),
|
|
)
|
|
loader = DiffusersPipelineLoader(LoadConfig(), cfg)
|
|
loader._hwr_state = None
|
|
|
|
model = nn.Module()
|
|
model.load_lora_weights = MagicMock()
|
|
|
|
loader._maybe_fuse_distilled_lora(model)
|
|
|
|
model.load_lora_weights.assert_called_once_with("/path/to/lora.safetensors")
|
|
assert getattr(model, "lora_is_fused", False) is True
|
|
|
|
|
|
def test_hwr_commit_failure_discards_model_and_reloads_without_hwr_or_mmap(tmp_path: Path, monkeypatch):
|
|
from vllm_omni.diffusion.model_loader import diffusers_loader as loader_module
|
|
|
|
loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config(tmp_path, tmp_path / "store"))
|
|
models: list[_DummyPipelineModel] = []
|
|
|
|
def init_model(*args, **kwargs):
|
|
del args, kwargs
|
|
model = _DummyPipelineModel(source_prefix="transformer.")
|
|
models.append(model)
|
|
return model
|
|
|
|
def commit_error(*args, **kwargs):
|
|
del args, kwargs
|
|
raise loader_module._HWRCommitError("restore commit failed")
|
|
|
|
monkeypatch.setattr(loader, "_init_from_load_format", init_model)
|
|
monkeypatch.setattr(loader, "_get_weight_sources", lambda _model: ())
|
|
monkeypatch.setattr(loader, "_resolve_hwr", commit_error)
|
|
monkeypatch.setattr(loader, "load_weights", lambda *args, **kwargs: None)
|
|
monkeypatch.setattr(loader, "_process_weights_after_loading", lambda *args, **kwargs: None)
|
|
monkeypatch.setattr(loader, "_apply_skip_softmax_calibration", lambda *args, **kwargs: None)
|
|
|
|
recovered = loader.load_model(load_device="cpu", device=torch.device("cpu"))
|
|
|
|
assert len(models) == 2
|
|
assert recovered is models[1]
|
|
assert loader.take_host_weight_plan() is None
|
|
|
|
|
|
def test_required_hwr_miss_fails_before_ordinary_loading_or_publication(
|
|
tmp_path: Path,
|
|
monkeypatch,
|
|
):
|
|
canonical_root = tmp_path / "canonical"
|
|
canonical_root.mkdir()
|
|
save_file(
|
|
{"weight": torch.arange(4, dtype=torch.float32).to(torch.bfloat16).reshape(2, 2)},
|
|
str(canonical_root / "model.safetensors"),
|
|
)
|
|
loader = DiffusersPipelineLoader(
|
|
LoadConfig(),
|
|
_hwr_config(canonical_root, tmp_path / "empty-store", mode="required"),
|
|
)
|
|
pipeline = _HWRPipeline(canonical_root)
|
|
monkeypatch.setattr(loader, "_init_from_load_format", lambda *args, **kwargs: pipeline)
|
|
|
|
with pytest.raises(RuntimeError, match="Host Weight Runtime resolution failed"):
|
|
loader.load_model(load_device="cpu", device=torch.device("cpu"))
|
|
|
|
assert pipeline.load_count == 0
|
|
assert loader.take_host_weight_plan() is None
|
|
|
|
|
|
def _make_dlo_online_quant_config(dp_size: int = 2) -> OmniDiffusionConfig:
|
|
return OmniDiffusionConfig(
|
|
model="",
|
|
dtype=torch.float32,
|
|
quantization_config="fp8",
|
|
parallel_config=DiffusionParallelConfig(
|
|
data_parallel_size=dp_size,
|
|
sequence_parallel_size=1,
|
|
),
|
|
enable_distributed_layerwise_offload=True,
|
|
dlo_use_allgather=True,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("dist_offload", "use_allgather", "mode"),
|
|
[
|
|
(False, False, "preferred"),
|
|
(True, True, "preferred"),
|
|
(True, False, "disabled"),
|
|
],
|
|
)
|
|
def test_hwr_disabled_for_noneligible_dlo_paths_without_store_interaction(
|
|
monkeypatch,
|
|
tmp_path,
|
|
dist_offload,
|
|
use_allgather,
|
|
mode,
|
|
):
|
|
"""Disabled and AllGather paths must never construct or probe HWR."""
|
|
from vllm_omni.host_weight_runtime import HostWeightRuntime
|
|
|
|
root = tmp_path / "must-not-be-touched"
|
|
loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config("dummy-model", root, mode=mode))
|
|
model = _DummyPipelineModel(source_prefix="transformer.")
|
|
modules = SimpleNamespace(dit_names=("transformer",), dits=(model.transformer,))
|
|
|
|
def unexpected_store_construction(*args, **kwargs):
|
|
raise AssertionError(f"HWR store interaction was not eligible: {args}, {kwargs}")
|
|
|
|
monkeypatch.setattr(HostWeightRuntime, "from_config", unexpected_store_construction)
|
|
assert (
|
|
loader._resolve_hwr(
|
|
model,
|
|
modules,
|
|
dist_offload=dist_offload,
|
|
use_allgather=use_allgather,
|
|
load_format="default",
|
|
sources=tuple(model.weights_sources),
|
|
)
|
|
is None
|
|
)
|
|
assert not root.exists()
|
|
|
|
|
|
def test_required_hwr_rejects_a_model_without_a_restore_contract(tmp_path):
|
|
loader = DiffusersPipelineLoader(
|
|
LoadConfig(),
|
|
_hwr_config("dummy-model", tmp_path / "store", mode="required"),
|
|
)
|
|
model = _DummyPipelineModel(source_prefix="transformer.")
|
|
modules = SimpleNamespace(dit_names=("transformer",), dits=(model.transformer,))
|
|
|
|
with pytest.raises(ValueError, match="restore contract"):
|
|
loader._resolve_hwr(
|
|
model,
|
|
modules,
|
|
dist_offload=True,
|
|
use_allgather=False,
|
|
load_format="default",
|
|
sources=tuple(model.weights_sources),
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("offline", [False, True])
|
|
def test_prepare_weights_honors_component_index_and_explicit_override(tmp_path, mocker, offline):
|
|
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
|
|
|
|
snapshot = tmp_path / "snapshot"
|
|
transformer = snapshot / "transformer"
|
|
transformer.mkdir(parents=True)
|
|
indexed_files = [f"diffusion_pytorch_model-{index:05d}-of-00002.safetensors" for index in (1, 2)]
|
|
stale_file = "diffusion_pytorch_model-00001-of-00008.safetensors"
|
|
index_path = transformer / "diffusion_pytorch_model.safetensors.index.json"
|
|
index_path.write_text(
|
|
json.dumps({"weight_map": {f"weight.{index}": filename for index, filename in enumerate(indexed_files)}})
|
|
)
|
|
for filename in indexed_files + [stale_file]:
|
|
(transformer / filename).touch()
|
|
|
|
mocker.patch.object(loader_mod.huggingface_hub.constants, "HF_HUB_OFFLINE", offline)
|
|
|
|
def download_index(*, filename, **_kwargs):
|
|
if filename == "transformer/diffusion_pytorch_model.safetensors.index.json":
|
|
return str(index_path)
|
|
raise loader_mod.huggingface_hub.errors.EntryNotFoundError(filename)
|
|
|
|
hub_api = mocker.Mock()
|
|
hub_api.hf_hub_download.side_effect = download_index
|
|
mocker.patch.object(loader_mod, "hf_api", return_value=hub_api)
|
|
indexed_download = mocker.patch.object(loader_mod, "download_weights_from_hf_specific", return_value=str(snapshot))
|
|
generic_download = mocker.patch.object(loader_mod, "download_weights_from_hf", return_value=str(snapshot))
|
|
|
|
loader = _make_loader_with_weights([])
|
|
cache_dir = str(tmp_path / "cache")
|
|
loader.load_config.download_dir = cache_dir
|
|
folder, files, use_safetensors = loader._prepare_weights(
|
|
"org/model",
|
|
subfolder="transformer",
|
|
revision="revision",
|
|
fall_back_to_pt=True,
|
|
allow_patterns_overrides=None,
|
|
)
|
|
|
|
assert folder == str(transformer)
|
|
assert [str(transformer / filename) for filename in indexed_files] == files
|
|
assert use_safetensors
|
|
assert hub_api.hf_hub_download.call_count == len(loader_mod.SAFETENSORS_INDEX_FILES)
|
|
hub_api.hf_hub_download.assert_any_call(
|
|
repo_id="org/model",
|
|
filename="transformer/diffusion_pytorch_model.safetensors.index.json",
|
|
cache_dir=cache_dir,
|
|
revision="revision",
|
|
local_files_only=offline,
|
|
)
|
|
indexed_download.assert_called_once_with(
|
|
model_name_or_path="org/model",
|
|
cache_dir=cache_dir,
|
|
allow_patterns=[f"transformer/{filename}" for filename in indexed_files],
|
|
revision="revision",
|
|
ignore_patterns=loader.load_config.ignore_patterns,
|
|
require_all=True,
|
|
)
|
|
|
|
_, override_files, _ = loader._prepare_weights(
|
|
"org/model",
|
|
subfolder="transformer",
|
|
revision="revision",
|
|
fall_back_to_pt=True,
|
|
allow_patterns_overrides=[stale_file],
|
|
)
|
|
assert override_files == [str(transformer / stale_file)]
|
|
generic_download.assert_called_once_with(
|
|
"org/model",
|
|
cache_dir,
|
|
[stale_file],
|
|
"revision",
|
|
subfolder="transformer",
|
|
ignore_patterns=loader.load_config.ignore_patterns,
|
|
)
|
|
|
|
|
|
def test_prepare_local_weights_honors_component_index(tmp_path):
|
|
transformer = tmp_path / "transformer"
|
|
transformer.mkdir()
|
|
indexed_file = "diffusion_pytorch_model-00001-of-00001.safetensors"
|
|
stale_file = "diffusion_pytorch_model-00001-of-00008.safetensors"
|
|
(transformer / "diffusion_pytorch_model.safetensors.index.json").write_text(
|
|
json.dumps({"weight_map": {"weight": indexed_file}})
|
|
)
|
|
for filename in (indexed_file, stale_file):
|
|
(transformer / filename).touch()
|
|
|
|
folder, files, use_safetensors = _make_loader_with_weights([])._prepare_weights(
|
|
tmp_path,
|
|
subfolder="transformer",
|
|
revision=None,
|
|
fall_back_to_pt=True,
|
|
allow_patterns_overrides=None,
|
|
)
|
|
|
|
assert folder == str(transformer)
|
|
assert files == [str(transformer / indexed_file)]
|
|
assert use_safetensors
|
|
|
|
|
|
def test_prepare_local_bin_index_is_authoritative_and_safetensors_remains_preferred(tmp_path):
|
|
transformer = tmp_path / "transformer"
|
|
transformer.mkdir()
|
|
indexed_files = [f"diffusion_pytorch_model-{index:05d}-of-00002.bin" for index in (1, 2)]
|
|
stale_file = "diffusion_pytorch_model-00001-of-00008.bin"
|
|
(transformer / "diffusion_pytorch_model.bin.index.json").write_text(
|
|
json.dumps({"weight_map": {f"weight.{index}": filename for index, filename in enumerate(indexed_files)}})
|
|
)
|
|
for filename in (*indexed_files, stale_file):
|
|
(transformer / filename).touch()
|
|
|
|
_, files, use_safetensors = _make_loader_with_weights([])._prepare_weights(
|
|
tmp_path,
|
|
subfolder="transformer",
|
|
revision=None,
|
|
fall_back_to_pt=True,
|
|
allow_patterns_overrides=None,
|
|
)
|
|
|
|
assert files == [str(transformer / filename) for filename in indexed_files]
|
|
assert not use_safetensors
|
|
|
|
safetensors_file = "diffusion_pytorch_model.safetensors"
|
|
(transformer / "diffusion_pytorch_model.safetensors.index.json").write_text(
|
|
json.dumps({"weight_map": {"weight": safetensors_file}})
|
|
)
|
|
(transformer / safetensors_file).touch()
|
|
_, files, use_safetensors = _make_loader_with_weights([])._prepare_weights(
|
|
tmp_path,
|
|
subfolder="transformer",
|
|
revision=None,
|
|
fall_back_to_pt=True,
|
|
allow_patterns_overrides=None,
|
|
)
|
|
|
|
assert files == [str(transformer / safetensors_file)]
|
|
assert use_safetensors
|
|
|
|
|
|
def test_prepare_weights_rejects_polluted_offline_cache_without_index(tmp_path, mocker):
|
|
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
|
|
|
|
snapshot = tmp_path / "snapshot"
|
|
transformer = snapshot / "transformer"
|
|
transformer.mkdir(parents=True)
|
|
for shard_count in (4, 8):
|
|
for shard in range(1, shard_count + 1):
|
|
(transformer / f"diffusion_pytorch_model-{shard:05d}-of-{shard_count:05d}.safetensors").touch()
|
|
|
|
mocker.patch.object(loader_mod.huggingface_hub.constants, "HF_HUB_OFFLINE", True)
|
|
hub_api = mocker.Mock()
|
|
hub_api.hf_hub_download.side_effect = loader_mod.huggingface_hub.errors.EntryNotFoundError("index is not cached")
|
|
mocker.patch.object(loader_mod, "hf_api", return_value=hub_api)
|
|
mocker.patch.object(loader_mod, "download_weights_from_hf", return_value=str(snapshot))
|
|
|
|
loader = _make_loader_with_weights([])
|
|
with pytest.raises(ValueError, match="conflicting shard totals"):
|
|
loader._prepare_weights(
|
|
"org/model",
|
|
subfolder="transformer",
|
|
revision="revision",
|
|
fall_back_to_pt=True,
|
|
allow_patterns_overrides=None,
|
|
)
|
|
|
|
assert hub_api.hf_hub_download.call_count == len(loader_mod.SAFETENSORS_INDEX_FILES) + len(
|
|
loader_mod.PT_INDEX_FILES
|
|
)
|
|
assert all(call.kwargs["local_files_only"] for call in hub_api.hf_hub_download.call_args_list)
|
|
|
|
|
|
def test_strict_check_only_validates_source_prefix_parameters():
|
|
model = _DummyPipelineModel(source_prefix="transformer.")
|
|
loader = _make_loader_with_weights(["transformer.weight"])
|
|
|
|
# Should not require VAE parameters because they are outside weights_sources.
|
|
loader.load_weights(model)
|
|
|
|
|
|
def test_strict_check_raises_when_source_parameters_are_missing():
|
|
model = _DummyPipelineModel(source_prefix="transformer.")
|
|
loader = _make_loader_with_weights([])
|
|
|
|
with pytest.raises(ValueError, match="transformer.weight"):
|
|
loader.load_weights(model)
|
|
|
|
|
|
def test_empty_source_prefix_keeps_full_model_strict_check():
|
|
model = _DummyPipelineModel(source_prefix="")
|
|
loader = _make_loader_with_weights(["transformer.weight"])
|
|
|
|
with pytest.raises(ValueError, match="vae.weight"):
|
|
loader.load_weights(model)
|
|
|
|
|
|
def test_stream_online_quant_weights_offloads_layers_after_processing():
|
|
from vllm.model_executor.model_loader.reload.layerwise import (
|
|
get_layerwise_info,
|
|
)
|
|
|
|
events: list[str] = []
|
|
|
|
class _OnlineQuantMethod:
|
|
uses_meta_device = True
|
|
|
|
class _TrackedLayer(nn.Linear):
|
|
def __init__(self, name: str):
|
|
super().__init__(2, 2, bias=False)
|
|
self.name = name
|
|
self.quant_method = _OnlineQuantMethod()
|
|
get_layerwise_info(self).load_numel_total = self.weight.numel()
|
|
|
|
def to(self, *args, **kwargs):
|
|
events.append(self.name)
|
|
return super().to(*args, **kwargs)
|
|
|
|
class _StreamingModel(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.first = _TrackedLayer("first")
|
|
self.second = _TrackedLayer("second")
|
|
|
|
model = _StreamingModel()
|
|
weights = iter(
|
|
[
|
|
("first.weight", torch.zeros((2, 2))),
|
|
("second.weight", torch.zeros((2, 2))),
|
|
]
|
|
)
|
|
streamed = DiffusersPipelineLoader._stream_online_quant_weights_to_cpu(model, weights)
|
|
|
|
assert next(streamed)[0] == "first.weight"
|
|
get_layerwise_info(model.first).reset()
|
|
assert next(streamed)[0] == "second.weight"
|
|
assert events == ["first"]
|
|
|
|
get_layerwise_info(model.second).reset()
|
|
with pytest.raises(StopIteration):
|
|
next(streamed)
|
|
assert events == ["first", "second"]
|
|
|
|
|
|
def test_process_weights_skips_completed_online_quant_layer(monkeypatch):
|
|
from unittest.mock import Mock
|
|
|
|
from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
|
|
from vllm.model_executor.model_loader.reload import layerwise
|
|
|
|
class _TrackedLayer(nn.Linear):
|
|
def __init__(self):
|
|
super().__init__(2, 2, bias=False)
|
|
self.to_calls: list[object] = []
|
|
|
|
def to(self, *args, **kwargs):
|
|
self.to_calls.append(args[0] if args else kwargs.get("device"))
|
|
return super().to(*args, **kwargs)
|
|
|
|
model = nn.Module()
|
|
model.layer = _TrackedLayer()
|
|
quant_method = Mock(spec=QuantizeMethodBase)
|
|
quant_method.uses_meta_device = True
|
|
model.layer.quant_method = quant_method
|
|
model.layer._already_called_process_weights_after_loading = True
|
|
finalize = Mock()
|
|
monkeypatch.setattr(layerwise, "finalize_layerwise_processing", finalize)
|
|
|
|
loader = _make_loader_with_weights([])
|
|
loader._process_weights_after_loading(model, torch.device("cuda"))
|
|
|
|
finalize.assert_called_once_with(model, model_config=None)
|
|
quant_method.process_weights_after_loading.assert_not_called()
|
|
assert model.layer.to_calls == []
|
|
|
|
|
|
class _ConfigAwareModel(nn.Module):
|
|
def __init__(self, *, od_config):
|
|
super().__init__()
|
|
self.captured_config = get_current_diffusion_config()
|
|
self.seen_config_during_init = get_current_diffusion_config_or_none()
|
|
self.od_config = od_config
|
|
|
|
|
|
def test_initialize_model_sets_current_diffusion_config_during_model_construction(monkeypatch):
|
|
import vllm_omni.diffusion.registry as registry_mod
|
|
|
|
od_config = SimpleNamespace(
|
|
model_class_name="DummyPipeline",
|
|
parallel_config=SimpleNamespace(vae_patch_parallel_size=1, sequence_parallel_size=1),
|
|
vae_use_slicing=False,
|
|
vae_use_tiling=False,
|
|
)
|
|
|
|
monkeypatch.setattr(
|
|
registry_mod.DiffusionModelRegistry,
|
|
"_try_load_model_cls",
|
|
staticmethod(lambda _name: _ConfigAwareModel),
|
|
)
|
|
monkeypatch.setattr(registry_mod, "_apply_sequence_parallel_if_enabled", lambda *_args, **_kwargs: None)
|
|
|
|
model = initialize_model(od_config)
|
|
|
|
assert model.captured_config is od_config
|
|
assert model.seen_config_during_init is od_config
|
|
assert get_current_diffusion_config_or_none() is None
|
|
|
|
|
|
def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch):
|
|
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
|
|
|
|
class _DeviceContext:
|
|
def __init__(self, device_type: str):
|
|
self.type = device_type
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc, tb):
|
|
return False
|
|
|
|
od_config = SimpleNamespace(
|
|
dtype=torch.float32,
|
|
parallel_config=SimpleNamespace(use_hsdp=False),
|
|
quantization_config=None,
|
|
)
|
|
|
|
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
|
|
loader.load_weights = lambda model: None # type: ignore[assignment]
|
|
loader._process_weights_after_loading = lambda model, target_device: None # type: ignore[assignment]
|
|
|
|
monkeypatch.setattr(loader_mod, "resolve_obj_by_qualname", lambda _name: _ConfigAwareModel)
|
|
monkeypatch.setattr(loader_mod.torch, "device", lambda _name: _DeviceContext("cpu"))
|
|
|
|
model = loader.load_model(
|
|
load_device="cpu",
|
|
load_format="custom_pipeline",
|
|
custom_pipeline_name="tests.dummy.ConfigAwarePipeline",
|
|
)
|
|
|
|
assert model.captured_config is od_config
|
|
assert model.seen_config_during_init is od_config
|
|
assert get_current_diffusion_config_or_none() is None
|
|
|
|
|
|
def test_dlo_transfers_loader_plan_and_skips_ordinary_weight_loading(monkeypatch):
|
|
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
|
|
|
|
od_config = SimpleNamespace(
|
|
dtype=torch.float32,
|
|
parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1),
|
|
quantization_config=None,
|
|
enable_distributed_layerwise_offload=True,
|
|
dlo_use_allgather=False,
|
|
model="unused",
|
|
)
|
|
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
|
|
model = nn.Module()
|
|
model.transformer = nn.Linear(2, 2, bias=False)
|
|
plan = HostWeightPlan(
|
|
backing_kind="checkpoint_mmap",
|
|
bindings={},
|
|
)
|
|
loaded_ordinary_weights = False
|
|
|
|
def load_weights(_model):
|
|
nonlocal loaded_ordinary_weights
|
|
loaded_ordinary_weights = True
|
|
|
|
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
|
|
loader.load_weights = load_weights # type: ignore[method-assign]
|
|
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
|
|
monkeypatch.setattr(
|
|
loader_mod,
|
|
"build_checkpoint_mmap_plan",
|
|
lambda *_args, **_kwargs: HostWeightPlanResult(plan),
|
|
)
|
|
|
|
assert loader.load_model(load_device="cpu") is model
|
|
assert not loaded_ordinary_weights
|
|
assert loader.take_host_weight_plan() is plan
|
|
assert loader.take_host_weight_plan() is None
|
|
|
|
|
|
def test_compact_rank_local_layer_offload_skips_dlo_mmap_planning(monkeypatch):
|
|
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
|
|
from vllm_omni.diffusion.offloader.config import materialize_legacy_offload_flags
|
|
|
|
od_config = SimpleNamespace(
|
|
dtype=torch.float32,
|
|
parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1),
|
|
quantization_config=None,
|
|
diffusion_offload_config={"mode": "layer", "components": ["dit"]},
|
|
enable_cpu_offload=False,
|
|
enable_layerwise_offload=False,
|
|
enable_distributed_layerwise_offload=False,
|
|
dlo_use_allgather=True,
|
|
dlo_resident_layers=0,
|
|
pin_cpu_memory=True,
|
|
model="unused",
|
|
)
|
|
materialize_legacy_offload_flags(od_config)
|
|
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
|
|
model = nn.Module()
|
|
model.transformer = nn.Linear(2, 2, bias=False)
|
|
calls: list[str] = []
|
|
|
|
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
|
|
loader.load_weights = lambda _model: calls.append("load") # type: ignore[method-assign]
|
|
loader._process_weights_after_loading = lambda *_args: calls.append("process") # type: ignore[method-assign]
|
|
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
|
|
monkeypatch.setattr(
|
|
loader_mod,
|
|
"build_checkpoint_mmap_plan",
|
|
lambda *_args, **_kwargs: pytest.fail("ordinary rank-local offload must not plan DLO mmap"),
|
|
)
|
|
|
|
assert loader.load_model(load_device="cpu") is model
|
|
assert calls == ["load", "process"]
|
|
assert loader.take_host_weight_plan() is None
|
|
|
|
|
|
class _UnsupportedOnlineQuantMethod:
|
|
uses_meta_device = True
|
|
|
|
|
|
def _compact_layer_offload_config(
|
|
components: list[str],
|
|
layer_options: dict[str, dict[str, str]],
|
|
) -> SimpleNamespace:
|
|
return SimpleNamespace(
|
|
dtype=torch.float32,
|
|
parallel_config=SimpleNamespace(
|
|
use_hsdp=False,
|
|
tensor_parallel_size=1,
|
|
data_parallel_size=2,
|
|
sequence_parallel_size=1,
|
|
),
|
|
quantization_config=None,
|
|
diffusion_offload_config={
|
|
"mode": "layer",
|
|
"components": components,
|
|
"layer_options": layer_options,
|
|
},
|
|
enable_cpu_offload=False,
|
|
enable_layerwise_offload=False,
|
|
enable_distributed_layerwise_offload=False,
|
|
dlo_use_allgather=True,
|
|
dlo_resident_layers=0,
|
|
pin_cpu_memory=False,
|
|
host_weight_runtime_mode="disabled",
|
|
model="unused",
|
|
)
|
|
|
|
|
|
def _stub_ordinary_loader(loader, model, monkeypatch, loader_mod) -> None:
|
|
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
|
|
loader.load_weights = lambda _model: None # type: ignore[method-assign]
|
|
loader._process_weights_after_loading = lambda *_args: None # type: ignore[method-assign]
|
|
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
|
|
monkeypatch.setattr(
|
|
loader_mod,
|
|
"build_checkpoint_mmap_plan",
|
|
lambda *_args, **_kwargs: HostWeightPlanResult(None, "test fallback"),
|
|
)
|
|
|
|
|
|
def test_text_encoder_allgather_rejects_unsupported_online_quantization(monkeypatch):
|
|
config = _compact_layer_offload_config(
|
|
["text_encoder"],
|
|
{"text_encoder": {"weight_transfer": "allgather"}},
|
|
)
|
|
loader = DiffusersPipelineLoader(LoadConfig(), config)
|
|
model = nn.Module()
|
|
model.text_encoder = nn.Linear(2, 2, bias=False)
|
|
model.text_encoder.quant_method = _UnsupportedOnlineQuantMethod()
|
|
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
|
|
|
|
with pytest.raises(ValueError, match="unsupported online methods: _UnsupportedOnlineQuantMethod"):
|
|
loader.load_model(load_device="cpu")
|
|
|
|
|
|
def test_dit_allgather_ignores_unsupported_quantization_on_unselected_encoder(monkeypatch):
|
|
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
|
|
|
|
config = _compact_layer_offload_config(
|
|
["dit"],
|
|
{"dit": {"weight_transfer": "allgather"}},
|
|
)
|
|
loader = DiffusersPipelineLoader(LoadConfig(), config)
|
|
model = nn.Module()
|
|
model.transformer = nn.Linear(2, 2, bias=False)
|
|
model.text_encoder = nn.Linear(2, 2, bias=False)
|
|
model.text_encoder.quant_method = _UnsupportedOnlineQuantMethod()
|
|
_stub_ordinary_loader(loader, model, monkeypatch, loader_mod)
|
|
|
|
assert loader.load_model(load_device="cpu") is model
|
|
|
|
|
|
def test_rank_local_encoder_quantization_does_not_mark_dit_mmap_plan_online(monkeypatch):
|
|
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
|
|
|
|
config = _compact_layer_offload_config(
|
|
["dit", "text_encoder"],
|
|
{
|
|
"dit": {"weight_transfer": "allgather"},
|
|
"text_encoder": {"weight_transfer": "rank-local"},
|
|
},
|
|
)
|
|
loader = DiffusersPipelineLoader(LoadConfig(), config)
|
|
model = nn.Module()
|
|
model.transformer = nn.Linear(2, 2, bias=False)
|
|
model.text_encoder = nn.Linear(2, 2, bias=False)
|
|
model.text_encoder.quant_method = _UnsupportedOnlineQuantMethod()
|
|
planned_online_quantization: list[bool] = []
|
|
|
|
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
|
|
loader.load_weights = lambda _model: None # type: ignore[method-assign]
|
|
loader._process_weights_after_loading = lambda *_args: None # type: ignore[method-assign]
|
|
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
|
|
|
|
def capture_plan(*_args, **kwargs):
|
|
planned_online_quantization.append(kwargs["online_quantization"])
|
|
return HostWeightPlanResult(None, "test fallback")
|
|
|
|
monkeypatch.setattr(loader_mod, "build_checkpoint_mmap_plan", capture_plan)
|
|
|
|
assert loader.load_model(load_device="cpu") is model
|
|
assert planned_online_quantization == [False]
|
|
|
|
|
|
def test_dlo_plan_loads_component_sources_outside_planned_dit(monkeypatch):
|
|
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
|
|
|
|
od_config = SimpleNamespace(
|
|
dtype=torch.float32,
|
|
parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1),
|
|
quantization_config=None,
|
|
enable_distributed_layerwise_offload=True,
|
|
dlo_use_allgather=False,
|
|
model="unused",
|
|
)
|
|
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
|
|
|
|
class MixedSourceModel(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.transformer = nn.Linear(2, 2, bias=False)
|
|
self.text_encoder = nn.Linear(2, 2, bias=False)
|
|
self.weights_sources = (
|
|
DiffusersPipelineLoader.ComponentSource("unused", None, None, prefix="transformer."),
|
|
DiffusersPipelineLoader.ComponentSource("unused", None, None, prefix="text_encoder."),
|
|
)
|
|
self.loaded_weight_names: list[str] = []
|
|
|
|
def load_weights(self, weights):
|
|
self.loaded_weight_names = [name for name, _ in weights]
|
|
return set(self.loaded_weight_names)
|
|
|
|
model = MixedSourceModel()
|
|
plan = HostWeightPlan(
|
|
backing_kind="checkpoint_mmap",
|
|
bindings={
|
|
"transformer.weight": TensorBinding(
|
|
checkpoint_key="weight",
|
|
file_path="unused",
|
|
)
|
|
},
|
|
planned_source_prefixes=frozenset({"transformer."}),
|
|
)
|
|
requested_prefixes: list[str] = []
|
|
|
|
def get_weights(source, model=None):
|
|
del model
|
|
requested_prefixes.append(source.prefix)
|
|
yield source.prefix + "weight", torch.ones(2, 2)
|
|
|
|
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
|
|
loader._get_weights_iterator = get_weights # type: ignore[method-assign]
|
|
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
|
|
monkeypatch.setattr(
|
|
loader_mod,
|
|
"build_checkpoint_mmap_plan",
|
|
lambda *_args, **_kwargs: HostWeightPlanResult(plan),
|
|
)
|
|
|
|
assert loader.load_model(load_device="cpu") is model
|
|
assert requested_prefixes == ["text_encoder."]
|
|
assert model.loaded_weight_names == ["text_encoder.weight"]
|
|
assert loader.take_host_weight_plan() is plan
|
|
|
|
|
|
def test_dlo_plan_fallback_runs_ordinary_loader(monkeypatch):
|
|
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
|
|
|
|
od_config = SimpleNamespace(
|
|
dtype=torch.float32,
|
|
parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1),
|
|
quantization_config=None,
|
|
enable_distributed_layerwise_offload=True,
|
|
dlo_use_allgather=False,
|
|
model="unused",
|
|
)
|
|
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
|
|
model = nn.Module()
|
|
model.transformer = nn.Linear(2, 2, bias=False)
|
|
calls: list[str] = []
|
|
|
|
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
|
|
loader.load_weights = lambda _model: calls.append("load") # type: ignore[method-assign]
|
|
loader._process_weights_after_loading = lambda *_args: calls.append("process") # type: ignore[method-assign]
|
|
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
|
|
monkeypatch.setattr(
|
|
loader_mod,
|
|
"build_checkpoint_mmap_plan",
|
|
lambda *_args, **_kwargs: HostWeightPlanResult(None, "not direct-compatible"),
|
|
)
|
|
|
|
assert loader.load_model(load_device="cpu") is model
|
|
assert calls == ["load", "process"]
|
|
assert loader.take_host_weight_plan() is None
|
|
|
|
|
|
def test_dlo_allgather_online_fp8_uses_ordinary_loader(monkeypatch):
|
|
from vllm.model_executor.layers.quantization.online.fp8 import (
|
|
Fp8PerTensorOnlineLinearMethod,
|
|
)
|
|
|
|
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
|
|
|
|
od_config = _make_dlo_online_quant_config()
|
|
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
|
|
model = nn.Module()
|
|
model.transformer = nn.Linear(2, 2, bias=False)
|
|
model.transformer.quant_method = object.__new__(Fp8PerTensorOnlineLinearMethod)
|
|
model.transformer.quant_method.uses_meta_device = True
|
|
calls: list[object] = []
|
|
allowlist_models: list[nn.Module] = []
|
|
|
|
original_allowlist_check = loader._unsupported_dlo_allgather_online_quant_methods
|
|
|
|
def check_allowlist(candidate: nn.Module) -> tuple[str, ...]:
|
|
allowlist_models.append(candidate)
|
|
return original_allowlist_check(candidate)
|
|
|
|
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
|
|
monkeypatch.setattr(loader, "_unsupported_dlo_allgather_online_quant_methods", check_allowlist)
|
|
loader._request_offload_after_quant = lambda _model: 1 # type: ignore[method-assign]
|
|
loader.load_weights = ( # type: ignore[method-assign]
|
|
lambda _model, *, stream_online_quant_to_cpu=False: calls.append(("load", stream_online_quant_to_cpu))
|
|
)
|
|
loader._process_weights_after_loading = lambda *_args: calls.append("process") # type: ignore[method-assign]
|
|
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
|
|
monkeypatch.setattr(
|
|
loader_mod,
|
|
"build_checkpoint_mmap_plan",
|
|
lambda *_args, **_kwargs: HostWeightPlanResult(
|
|
None,
|
|
"online quantization requires the ordinary loader",
|
|
),
|
|
)
|
|
|
|
assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model
|
|
assert allowlist_models == [model.transformer]
|
|
assert calls == [("load", True), "process"]
|
|
assert loader.take_host_weight_plan() is None
|
|
|
|
|
|
def test_dlo_allgather_online_int8_uses_ordinary_loader(monkeypatch):
|
|
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
|
|
from vllm_omni.quantization.int8_config import NPUInt8OnlineLinearMethod
|
|
|
|
od_config = _make_dlo_online_quant_config()
|
|
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
|
|
model = nn.Module()
|
|
model.transformer = nn.Linear(2, 2, bias=False)
|
|
model.transformer.quant_method = object.__new__(NPUInt8OnlineLinearMethod)
|
|
model.transformer.quant_method.uses_meta_device = True
|
|
calls: list[object] = []
|
|
allowlist_models: list[nn.Module] = []
|
|
|
|
original_allowlist_check = loader._unsupported_dlo_allgather_online_quant_methods
|
|
|
|
def check_allowlist(candidate: nn.Module) -> tuple[str, ...]:
|
|
allowlist_models.append(candidate)
|
|
return original_allowlist_check(candidate)
|
|
|
|
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
|
|
monkeypatch.setattr(loader, "_unsupported_dlo_allgather_online_quant_methods", check_allowlist)
|
|
loader._request_offload_after_quant = lambda _model: 1 # type: ignore[method-assign]
|
|
loader.load_weights = ( # type: ignore[method-assign]
|
|
lambda _model, *, stream_online_quant_to_cpu=False: calls.append(("load", stream_online_quant_to_cpu))
|
|
)
|
|
loader._process_weights_after_loading = lambda *_args: calls.append("process") # type: ignore[method-assign]
|
|
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
|
|
monkeypatch.setattr(
|
|
loader_mod,
|
|
"build_checkpoint_mmap_plan",
|
|
lambda *_args, **_kwargs: HostWeightPlanResult(
|
|
None,
|
|
"online quantization requires the ordinary loader",
|
|
),
|
|
)
|
|
|
|
assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model
|
|
assert allowlist_models == [model.transformer]
|
|
assert calls == [("load", True), "process"]
|
|
assert loader.take_host_weight_plan() is None
|
|
|
|
|
|
def test_dlo_online_quant_group_size_one_skips_allgather_gate(monkeypatch):
|
|
"""A DLO group of one runs no weight collective, so an otherwise
|
|
unvalidated online method must not be rejected by the AllGather gate."""
|
|
|
|
class UnsupportedOnlineMethod:
|
|
uses_meta_device = True
|
|
|
|
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
|
|
|
|
od_config = _make_dlo_online_quant_config(dp_size=1)
|
|
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
|
|
model = nn.Module()
|
|
model.transformer = nn.Linear(2, 2, bias=False)
|
|
model.transformer.quant_method = UnsupportedOnlineMethod()
|
|
calls: list[object] = []
|
|
allowlist_models: list[nn.Module] = []
|
|
|
|
original_allowlist_check = loader._unsupported_dlo_allgather_online_quant_methods
|
|
|
|
def check_allowlist(candidate: nn.Module) -> tuple[str, ...]:
|
|
allowlist_models.append(candidate)
|
|
return original_allowlist_check(candidate)
|
|
|
|
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
|
|
monkeypatch.setattr(loader, "_unsupported_dlo_allgather_online_quant_methods", check_allowlist)
|
|
loader._request_offload_after_quant = lambda _model: 1 # type: ignore[method-assign]
|
|
loader.load_weights = ( # type: ignore[method-assign]
|
|
lambda _model, *, stream_online_quant_to_cpu=False: calls.append(("load", stream_online_quant_to_cpu))
|
|
)
|
|
loader._process_weights_after_loading = lambda *_args: calls.append("process") # type: ignore[method-assign]
|
|
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
|
|
monkeypatch.setattr(
|
|
loader_mod,
|
|
"build_checkpoint_mmap_plan",
|
|
lambda *_args, **_kwargs: HostWeightPlanResult(
|
|
None,
|
|
"online quantization requires the ordinary loader",
|
|
),
|
|
)
|
|
|
|
assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model
|
|
assert allowlist_models == []
|
|
assert calls == [("load", True), "process"]
|
|
|
|
|
|
def test_dlo_allgather_rejects_unvalidated_online_quant_method(monkeypatch):
|
|
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
|
|
|
|
class UnsupportedOnlineMethod:
|
|
uses_meta_device = True
|
|
|
|
od_config = _make_dlo_online_quant_config()
|
|
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
|
|
model = nn.Module()
|
|
model.transformer = nn.Linear(2, 2, bias=False)
|
|
model.transformer.quant_method = UnsupportedOnlineMethod()
|
|
|
|
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
|
|
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
|
|
monkeypatch.setattr(
|
|
loader_mod,
|
|
"build_checkpoint_mmap_plan",
|
|
lambda *_args, **_kwargs: HostWeightPlanResult(
|
|
None,
|
|
"online quantization requires the ordinary loader",
|
|
),
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="per-tensor FP8, INT8, and MXFP8 linears"):
|
|
loader.load_model(load_device="cpu")
|
|
|
|
|
|
def test_dlo_allgather_online_mxfp8_uses_ordinary_loader(monkeypatch):
|
|
mxfp8_config = pytest.importorskip("vllm_omni.quantization.mxfp8_config")
|
|
NPUMxfp8OnlineLinearMethod = mxfp8_config.NPUMxfp8OnlineLinearMethod
|
|
|
|
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
|
|
|
|
od_config = _make_dlo_online_quant_config()
|
|
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
|
|
model = nn.Module()
|
|
model.transformer = nn.Linear(2, 2, bias=False)
|
|
model.transformer.quant_method = object.__new__(NPUMxfp8OnlineLinearMethod)
|
|
model.transformer.quant_method.uses_meta_device = True
|
|
calls: list[object] = []
|
|
allowlist_models: list[nn.Module] = []
|
|
|
|
original_allowlist_check = loader._unsupported_dlo_allgather_online_quant_methods
|
|
|
|
def check_allowlist(candidate: nn.Module) -> tuple[str, ...]:
|
|
allowlist_models.append(candidate)
|
|
return original_allowlist_check(candidate)
|
|
|
|
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
|
|
monkeypatch.setattr(loader, "_unsupported_dlo_allgather_online_quant_methods", check_allowlist)
|
|
loader._request_offload_after_quant = lambda _model: 1 # type: ignore[method-assign]
|
|
loader.load_weights = ( # type: ignore[method-assign]
|
|
lambda _model, *, stream_online_quant_to_cpu=False: calls.append(("load", stream_online_quant_to_cpu))
|
|
)
|
|
loader._process_weights_after_loading = lambda *_args: calls.append("process") # type: ignore[method-assign]
|
|
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
|
|
monkeypatch.setattr(
|
|
loader_mod,
|
|
"build_checkpoint_mmap_plan",
|
|
lambda *_args, **_kwargs: HostWeightPlanResult(
|
|
None,
|
|
"online quantization requires the ordinary loader",
|
|
),
|
|
)
|
|
|
|
assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model
|
|
assert allowlist_models == [model.transformer]
|
|
assert calls == [("load", True), "process"]
|
|
assert loader.take_host_weight_plan() is None
|
|
|
|
|
|
def test_hsdp_processes_quantized_weights_before_sharding(mocker):
|
|
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
|
|
from vllm_omni.diffusion.offloader.module_collector import PipelineModules
|
|
|
|
od_config = SimpleNamespace(
|
|
dtype=torch.float32,
|
|
parallel_config=SimpleNamespace(
|
|
use_hsdp=True,
|
|
hsdp_replicate_size=1,
|
|
hsdp_shard_size=2,
|
|
),
|
|
quantization_config=None,
|
|
)
|
|
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
|
|
loader.quant_config = object()
|
|
|
|
model = nn.Module()
|
|
model.transformer = nn.Linear(2, 2, bias=False)
|
|
events: list[str] = []
|
|
|
|
loader._init_from_load_format = mocker.Mock(return_value=model) # type: ignore[method-assign]
|
|
loader.load_weights = mocker.Mock(side_effect=lambda _model: events.append("load")) # type: ignore[method-assign]
|
|
loader._process_weights_after_loading = mocker.Mock( # type: ignore[method-assign]
|
|
side_effect=lambda _model, _device: events.append("process")
|
|
)
|
|
mocker.patch.object(
|
|
loader_mod.ModuleDiscovery,
|
|
"discover",
|
|
return_value=PipelineModules(
|
|
dits=[model.transformer],
|
|
dit_names=["transformer"],
|
|
vaes=[],
|
|
encoders=[],
|
|
encoder_names=[],
|
|
resident_modules=[],
|
|
resident_names=[],
|
|
),
|
|
)
|
|
mocker.patch(
|
|
"vllm_omni.diffusion.quantization.hsdp_fp8.prepare_fp8_layers_for_fsdp",
|
|
side_effect=lambda _model: events.append("prepare"),
|
|
)
|
|
mocker.patch.object(
|
|
loader_mod,
|
|
"apply_hsdp_to_model",
|
|
side_effect=lambda *_args, **_kwargs: events.append("shard"),
|
|
)
|
|
|
|
loader._load_model_with_hsdp(torch.device("cpu"))
|
|
|
|
assert events == ["load", "process", "prepare", "shard"]
|
|
|
|
|
|
def test_hsdp_fuses_distilled_lora_before_sharding(mocker):
|
|
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
|
|
from vllm_omni.diffusion.offloader.module_collector import PipelineModules
|
|
|
|
od_config = SimpleNamespace(
|
|
dtype=torch.float32,
|
|
parallel_config=SimpleNamespace(
|
|
use_hsdp=True,
|
|
hsdp_replicate_size=1,
|
|
hsdp_shard_size=2,
|
|
),
|
|
quantization_config=None,
|
|
lora_backend="distill",
|
|
lora_path="/path/to/lora.safetensors",
|
|
lora_scale=1.0,
|
|
)
|
|
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
|
|
loader.quant_config = None
|
|
|
|
model = nn.Module()
|
|
model.transformer = nn.Linear(2, 2, bias=False)
|
|
events: list[str] = []
|
|
|
|
model.load_lora_weights = mocker.Mock(side_effect=lambda path: events.append(f"fuse_lora:{path}"))
|
|
|
|
loader._init_from_load_format = mocker.Mock(return_value=model)
|
|
loader.load_weights = mocker.Mock(side_effect=lambda _model: events.append("load"))
|
|
loader._process_weights_after_loading = mocker.Mock(side_effect=lambda _model, _device: events.append("process"))
|
|
mocker.patch.object(
|
|
loader_mod.ModuleDiscovery,
|
|
"discover",
|
|
return_value=PipelineModules(
|
|
dits=[model.transformer],
|
|
dit_names=["transformer"],
|
|
vaes=[],
|
|
encoders=[],
|
|
encoder_names=[],
|
|
resident_modules=[],
|
|
resident_names=[],
|
|
),
|
|
)
|
|
mocker.patch.object(
|
|
loader_mod,
|
|
"apply_hsdp_to_model",
|
|
side_effect=lambda *_args, **_kwargs: events.append("shard"),
|
|
)
|
|
|
|
loader._load_model_with_hsdp(torch.device("cpu"))
|
|
|
|
assert events == ["load", "fuse_lora:/path/to/lora.safetensors", "process", "shard"]
|
|
assert getattr(model, "lora_is_fused", False) is True
|
|
|
|
|
|
def test_hsdp_fuses_multi_file_distilled_lora_wan22(mocker):
|
|
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
|
|
from vllm_omni.diffusion.offloader.module_collector import PipelineModules
|
|
|
|
od_config = SimpleNamespace(
|
|
dtype=torch.float32,
|
|
parallel_config=SimpleNamespace(
|
|
use_hsdp=True,
|
|
hsdp_replicate_size=1,
|
|
hsdp_shard_size=2,
|
|
),
|
|
quantization_config=None,
|
|
lora_backend="distill",
|
|
lora_path=["/path/to/high.safetensors", "/path/to/low.safetensors"],
|
|
lora_scale=1.0,
|
|
)
|
|
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
|
|
loader.quant_config = None
|
|
|
|
model = nn.Module()
|
|
model.transformer = nn.Linear(2, 2, bias=False)
|
|
model.transformer_2 = nn.Linear(2, 2, bias=False)
|
|
events: list[str] = []
|
|
|
|
model.load_lora_weights = mocker.Mock(side_effect=lambda paths: events.append(f"fuse_lora:{paths}"))
|
|
|
|
loader._init_from_load_format = mocker.Mock(return_value=model)
|
|
loader.load_weights = mocker.Mock(side_effect=lambda _model: events.append("load"))
|
|
loader._process_weights_after_loading = mocker.Mock(side_effect=lambda _model, _device: events.append("process"))
|
|
mocker.patch.object(
|
|
loader_mod.ModuleDiscovery,
|
|
"discover",
|
|
return_value=PipelineModules(
|
|
dits=[model.transformer, model.transformer_2],
|
|
dit_names=["transformer", "transformer_2"],
|
|
vaes=[],
|
|
encoders=[],
|
|
encoder_names=[],
|
|
resident_modules=[],
|
|
resident_names=[],
|
|
),
|
|
)
|
|
mocker.patch.object(
|
|
loader_mod,
|
|
"apply_hsdp_to_model",
|
|
side_effect=lambda *_args, **_kwargs: events.append("shard"),
|
|
)
|
|
|
|
loader._load_model_with_hsdp(torch.device("cpu"))
|
|
|
|
assert events == [
|
|
"load",
|
|
"fuse_lora:['/path/to/high.safetensors', '/path/to/low.safetensors']",
|
|
"process",
|
|
"shard",
|
|
"shard",
|
|
]
|
|
assert getattr(model, "lora_is_fused", False) is True
|
|
|
|
|
|
def test_load_model_fuses_distilled_lora_non_hsdp(mocker):
|
|
od_config = SimpleNamespace(
|
|
dtype=torch.float32,
|
|
parallel_config=SimpleNamespace(
|
|
use_hsdp=False,
|
|
tensor_parallel_size=1,
|
|
data_parallel_size=1,
|
|
sequence_parallel_size=1,
|
|
),
|
|
enable_cpu_offload=False,
|
|
quantization_config=None,
|
|
lora_backend="distill",
|
|
lora_path="/path/to/lora.safetensors",
|
|
lora_scale=1.0,
|
|
host_weight_runtime_mode="disabled",
|
|
)
|
|
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
|
|
loader._force_canonical_load = True
|
|
|
|
model = nn.Module()
|
|
model.transformer = nn.Linear(2, 2, bias=False)
|
|
events: list[str] = []
|
|
|
|
model.load_lora_weights = mocker.Mock(side_effect=lambda path: events.append(f"fuse_lora:{path}"))
|
|
|
|
loader._init_from_load_format = mocker.Mock(return_value=model)
|
|
loader.load_weights = mocker.Mock(side_effect=lambda _model: events.append("load"))
|
|
loader._process_weights_after_loading = mocker.Mock(side_effect=lambda _model, _device: events.append("process"))
|
|
loader._apply_skip_softmax_calibration = mocker.Mock()
|
|
|
|
res = loader.load_model(load_device="cpu", device=torch.device("cpu"))
|
|
|
|
assert res is model
|
|
assert events == ["load", "fuse_lora:/path/to/lora.safetensors", "process"]
|
|
assert getattr(model, "lora_is_fused", False) is True
|
|
|
|
|
|
def test_get_all_weights(prefetch_helios_model, mock_tp_group):
|
|
"""Ensure that get all weights on a tiny model resolves to nonempty weights."""
|
|
od_config = OmniDiffusionConfig(
|
|
model_class_name="HeliosPipeline",
|
|
model=model_path,
|
|
)
|
|
loader = DiffusersPipelineLoader(
|
|
load_config=LoadConfig(),
|
|
od_config=od_config,
|
|
)
|
|
pipeline = HeliosPipeline(od_config=od_config)
|
|
|
|
weights = list(loader.get_all_weights(pipeline))
|
|
assert len(weights) > 0
|
|
|
|
|
|
def test_load_model(prefetch_helios_model, mock_tp_group):
|
|
"""Ensure that load model creates an instance of the expected pipeline class."""
|
|
od_config = OmniDiffusionConfig(
|
|
model_class_name="HeliosPipeline",
|
|
model=model_path,
|
|
)
|
|
loader = DiffusersPipelineLoader(
|
|
load_config=LoadConfig(),
|
|
od_config=od_config,
|
|
)
|
|
model = loader.load_model(load_device="cpu")
|
|
assert isinstance(model, HeliosPipeline)
|