270 lines
11 KiB
Python
270 lines
11 KiB
Python
|
# SPDX-License-Identifier: Apache-2.0
|
||
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||
|
|
"""Exact projection reuse without model-specific schedules or artifacts."""
|
||
|
|
|
||
|
|
from types import SimpleNamespace
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
import torch
|
||
|
|
|
||
|
|
from vllm_omni.diffusion.cache.exact_projection_cache import ExactProjectionCache
|
||
|
|
|
||
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.fixture(autouse=True)
|
||
|
|
def tp1(monkeypatch):
|
||
|
|
import vllm.distributed
|
||
|
|
|
||
|
|
monkeypatch.setattr(vllm.distributed, "get_tensor_model_parallel_world_size", lambda: 1)
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.fixture(autouse=True)
|
||
|
|
def deterministic_weights():
|
||
|
|
with torch.random.fork_rng(devices=[]):
|
||
|
|
torch.manual_seed(0)
|
||
|
|
yield
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.parametrize("batch,in_features,out_features", [(2, 4, 8), (1, 8, 48), (4, 6, 18)])
|
||
|
|
def test_runtime_reuses_exact_outputs_and_tracks_changed_weights(mocker, batch, in_features, out_features):
|
||
|
|
cache = ExactProjectionCache()
|
||
|
|
linear = torch.nn.Linear(in_features, out_features)
|
||
|
|
x = torch.randn(batch, in_features)
|
||
|
|
compute = mocker.Mock(side_effect=lambda: linear(x))
|
||
|
|
with torch.no_grad():
|
||
|
|
cache.prepare(x)
|
||
|
|
first = cache.project("block", linear, x, compute)
|
||
|
|
first_copy = first.clone()
|
||
|
|
first.zero_() # The caller/provider cannot corrupt the owned cache.
|
||
|
|
assert torch.equal(cache.project("block", linear, x, compute), first_copy)
|
||
|
|
assert compute.call_count == 1
|
||
|
|
linear.weight.add_(0.25)
|
||
|
|
assert torch.equal(cache.project("block", linear, x, compute), linear(x))
|
||
|
|
assert compute.call_count == 2
|
||
|
|
other = x.clone()
|
||
|
|
other[0, 0] += 1
|
||
|
|
cache.prepare(other)
|
||
|
|
cache.project("block", linear, other, lambda: linear(other))
|
||
|
|
cache.prepare(x)
|
||
|
|
cache.project("block", linear, x, compute)
|
||
|
|
assert compute.call_count == 2
|
||
|
|
assert cache.hits == 2
|
||
|
|
|
||
|
|
|
||
|
|
def test_runtime_budget_retains_reusable_entries_without_schedule_thrashing(mocker):
|
||
|
|
cache = ExactProjectionCache(max_bytes=32)
|
||
|
|
linear = torch.nn.Linear(2, 4)
|
||
|
|
x = torch.ones(2, 2)
|
||
|
|
compute = mocker.Mock(side_effect=lambda: linear(x))
|
||
|
|
with torch.no_grad():
|
||
|
|
cache.prepare(x)
|
||
|
|
for name in ("first", "second", "first"):
|
||
|
|
cache.project(name, linear, x, compute)
|
||
|
|
assert cache._bytes <= 32
|
||
|
|
assert compute.call_count == 2
|
||
|
|
assert cache.hits == 1
|
||
|
|
cache.clear()
|
||
|
|
assert cache._bytes == 0 and not cache._entries
|
||
|
|
|
||
|
|
|
||
|
|
def test_dynamic_lora_buffers_and_suspend_mask_invalidate_entries(mocker):
|
||
|
|
cache = ExactProjectionCache()
|
||
|
|
linear = torch.nn.Linear(2, 2)
|
||
|
|
linear.lora_a_stacked = (torch.ones(1, 2),)
|
||
|
|
linear.lora_b_stacked = (torch.ones(2, 1),)
|
||
|
|
linear._diffusion_lora_active_slices = (True,)
|
||
|
|
x = torch.ones(1, 2)
|
||
|
|
|
||
|
|
def compute():
|
||
|
|
result = linear(x)
|
||
|
|
if linear._diffusion_lora_active_slices[0]:
|
||
|
|
result = result + x @ linear.lora_a_stacked[0].T @ linear.lora_b_stacked[0].T
|
||
|
|
return result
|
||
|
|
|
||
|
|
counted = mocker.Mock(side_effect=compute)
|
||
|
|
with torch.no_grad():
|
||
|
|
cache.prepare(x)
|
||
|
|
cache.project("a", linear, x, counted)
|
||
|
|
cache.project("a", linear, x, counted)
|
||
|
|
assert counted.call_count == 1
|
||
|
|
linear.lora_b_stacked[0].add_(1)
|
||
|
|
assert torch.equal(cache.project("a", linear, x, counted), compute())
|
||
|
|
linear._diffusion_lora_active_slices = (False,)
|
||
|
|
assert torch.equal(cache.project("a", linear, x, counted), compute())
|
||
|
|
assert counted.call_count == 3
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.parametrize("scope", ["local", "global", "child"])
|
||
|
|
@pytest.mark.parametrize("kind", ["pre", "post"])
|
||
|
|
def test_linear_hooks_run_even_after_cache_warmup(scope, kind):
|
||
|
|
cache = ExactProjectionCache()
|
||
|
|
linear = torch.nn.Linear(2, 2)
|
||
|
|
projection = torch.nn.Sequential(linear) if scope == "child" else linear
|
||
|
|
x = torch.ones(1, 2)
|
||
|
|
with torch.no_grad():
|
||
|
|
cache.prepare(x)
|
||
|
|
expected = projection(x)
|
||
|
|
cache.project("a", projection, x, lambda: projection(x))
|
||
|
|
calls = []
|
||
|
|
|
||
|
|
def pre_hook(module, args):
|
||
|
|
if module is linear:
|
||
|
|
calls.append(module)
|
||
|
|
return (args[0] + 1,)
|
||
|
|
|
||
|
|
def post_hook(module, args, output):
|
||
|
|
if module is linear:
|
||
|
|
calls.append(module)
|
||
|
|
return output + 1
|
||
|
|
|
||
|
|
if kind == "pre":
|
||
|
|
expected = linear(x + 1)
|
||
|
|
register = (
|
||
|
|
torch.nn.modules.module.register_module_forward_pre_hook
|
||
|
|
if scope == "global"
|
||
|
|
else linear.register_forward_pre_hook
|
||
|
|
)
|
||
|
|
hook = register(pre_hook)
|
||
|
|
else:
|
||
|
|
expected = expected + 1
|
||
|
|
register = (
|
||
|
|
torch.nn.modules.module.register_module_forward_hook
|
||
|
|
if scope == "global"
|
||
|
|
else linear.register_forward_hook
|
||
|
|
)
|
||
|
|
hook = register(post_hook)
|
||
|
|
try:
|
||
|
|
for _ in range(2):
|
||
|
|
assert torch.equal(cache.project("a", projection, x, lambda: projection(x)), expected)
|
||
|
|
assert len(calls) == 2
|
||
|
|
finally:
|
||
|
|
hook.remove()
|
||
|
|
assert torch.equal(cache.project("a", projection, x, lambda: projection(x)), projection(x))
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.parametrize("autocast_at_prepare", [False, True])
|
||
|
|
def test_numerical_settings_follow_projection_context(mocker, autocast_at_prepare):
|
||
|
|
cache = ExactProjectionCache()
|
||
|
|
linear = torch.nn.Linear(4, 4)
|
||
|
|
x = torch.randn(2, 4)
|
||
|
|
compute = mocker.Mock(side_effect=lambda: linear(x))
|
||
|
|
with torch.no_grad():
|
||
|
|
with torch.autocast("cpu", dtype=torch.bfloat16, enabled=autocast_at_prepare):
|
||
|
|
cache.prepare(x)
|
||
|
|
for enabled in (False, True, False, True):
|
||
|
|
with torch.autocast("cpu", dtype=torch.bfloat16, enabled=enabled):
|
||
|
|
actual = cache.project("a", linear, x, compute)
|
||
|
|
expected = linear(x)
|
||
|
|
assert actual.dtype == expected.dtype
|
||
|
|
assert torch.equal(actual, expected)
|
||
|
|
assert compute.call_count == 2
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.parametrize(
|
||
|
|
"setting",
|
||
|
|
["allow_tf32", "allow_bf16_reduced_precision_reduction", "allow_fp16_reduced_precision_reduction"],
|
||
|
|
)
|
||
|
|
def test_matmul_settings_invalidate_prepared_projection(monkeypatch, mocker, setting):
|
||
|
|
cache = ExactProjectionCache()
|
||
|
|
linear = torch.nn.Linear(2, 2)
|
||
|
|
x = torch.ones(1, 2)
|
||
|
|
compute = mocker.Mock(side_effect=lambda: linear(x))
|
||
|
|
with torch.no_grad():
|
||
|
|
cache.prepare(x)
|
||
|
|
cache.project("a", linear, x, compute)
|
||
|
|
monkeypatch.setattr(torch.backends.cuda.matmul, setting, not getattr(torch.backends.cuda.matmul, setting))
|
||
|
|
assert torch.equal(cache.project("a", linear, x, compute), linear(x))
|
||
|
|
assert compute.call_count == 2
|
||
|
|
|
||
|
|
|
||
|
|
def test_disabled_and_grad_paths_keep_original_computation(monkeypatch, mocker):
|
||
|
|
import vllm_omni.diffusion.cache.exact_projection_cache as module
|
||
|
|
|
||
|
|
monkeypatch.setattr(module, "tensor_digest", mocker.Mock(side_effect=AssertionError("must not hash")))
|
||
|
|
linear = torch.nn.Linear(2, 2)
|
||
|
|
x = torch.ones(1, 2, requires_grad=True)
|
||
|
|
for cache in (ExactProjectionCache(max_bytes=0), ExactProjectionCache()):
|
||
|
|
cache.prepare(x)
|
||
|
|
cache.project("a", linear, x, lambda: linear(x)).sum().backward()
|
||
|
|
assert x.grad is not None
|
||
|
|
with torch.no_grad():
|
||
|
|
cache = ExactProjectionCache(max_bytes=0)
|
||
|
|
cache.prepare(x)
|
||
|
|
assert torch.equal(cache.project("a", linear, x, lambda: linear(x)), linear(x))
|
||
|
|
|
||
|
|
|
||
|
|
def test_tp_rank_miss_forces_every_rank_to_recompute(monkeypatch, mocker):
|
||
|
|
import vllm.distributed
|
||
|
|
|
||
|
|
cache = ExactProjectionCache()
|
||
|
|
linear = torch.nn.Linear(2, 2)
|
||
|
|
x = torch.ones(1, 2)
|
||
|
|
compute = mocker.Mock(side_effect=lambda: linear(x))
|
||
|
|
with torch.no_grad():
|
||
|
|
cache.prepare(x)
|
||
|
|
cache.project("a", linear, x, compute)
|
||
|
|
votes = mocker.Mock(side_effect=[torch.tensor([1]), torch.tensor([2])])
|
||
|
|
monkeypatch.setattr(vllm.distributed, "get_tensor_model_parallel_world_size", lambda: 2)
|
||
|
|
monkeypatch.setattr(vllm.distributed, "get_tp_group", lambda: SimpleNamespace(all_reduce=votes))
|
||
|
|
cache.project("a", linear, x, compute)
|
||
|
|
assert compute.call_count == 2
|
||
|
|
cache.project("a", linear, x, compute)
|
||
|
|
assert compute.call_count == 2
|
||
|
|
assert votes.call_count == 2
|
||
|
|
|
||
|
|
|
||
|
|
def test_compiled_execution_bypasses_runtime_and_precomputed_results(monkeypatch, mocker):
|
||
|
|
import vllm_omni.diffusion.cache.exact_projection_cache as module
|
||
|
|
|
||
|
|
cache = ExactProjectionCache()
|
||
|
|
linear = torch.nn.Linear(2, 2)
|
||
|
|
x = torch.ones(1, 2)
|
||
|
|
compute = mocker.Mock(side_effect=lambda: linear(x))
|
||
|
|
with torch.no_grad():
|
||
|
|
cache.prepare(x)
|
||
|
|
cache.project("modulation", linear, x, compute)
|
||
|
|
monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True)
|
||
|
|
monkeypatch.setattr(module, "tensor_digest", mocker.Mock(side_effect=AssertionError("must not hash")))
|
||
|
|
seeded = mocker.patch.object(cache, "_lookup_precomputed", side_effect=AssertionError("must not seed"))
|
||
|
|
# An input prepared before compilation must also bypass the warm entry.
|
||
|
|
assert torch.equal(cache.project("modulation", linear, x, compute), linear(x))
|
||
|
|
cache.prepare(x)
|
||
|
|
assert torch.equal(cache.project("modulation", linear, x, compute), linear(x))
|
||
|
|
assert compute.call_count == 3
|
||
|
|
seeded.assert_not_called()
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.parametrize("world_size", [1, 2])
|
||
|
|
def test_precomputed_extension_is_used_only_at_tp1(monkeypatch, mocker, world_size):
|
||
|
|
import vllm.distributed
|
||
|
|
|
||
|
|
cache = ExactProjectionCache()
|
||
|
|
linear = torch.nn.Linear(2, 2)
|
||
|
|
x = torch.ones(1, 2)
|
||
|
|
monkeypatch.setattr(vllm.distributed, "get_tensor_model_parallel_world_size", lambda: world_size)
|
||
|
|
votes = mocker.Mock(side_effect=[torch.tensor([0]), torch.tensor([world_size])])
|
||
|
|
monkeypatch.setattr(vllm.distributed, "get_tp_group", lambda: SimpleNamespace(all_reduce=votes))
|
||
|
|
with torch.no_grad():
|
||
|
|
expected = linear(x)
|
||
|
|
seeded = mocker.patch.object(cache, "_lookup_precomputed", return_value=expected)
|
||
|
|
compute = mocker.Mock(side_effect=lambda: linear(x))
|
||
|
|
cache.prepare(x)
|
||
|
|
for _ in range(2):
|
||
|
|
assert torch.equal(cache.project("conditioning.scale_shift", linear, x, compute), expected)
|
||
|
|
assert seeded.call_count == (1 if world_size == 1 else 0)
|
||
|
|
assert compute.call_count == (0 if world_size == 1 else 1)
|
||
|
|
assert cache.hits == 1
|
||
|
|
|
||
|
|
|
||
|
|
def test_inference_weights_without_versions_use_safe_runtime_path(mocker):
|
||
|
|
with torch.inference_mode():
|
||
|
|
linear = torch.nn.Linear(2, 2)
|
||
|
|
x = torch.ones(1, 2)
|
||
|
|
cache = ExactProjectionCache()
|
||
|
|
compute = mocker.Mock(side_effect=lambda: linear(x))
|
||
|
|
cache.prepare(x)
|
||
|
|
cache.project("a", linear, x, compute)
|
||
|
|
linear.weight.add_(1)
|
||
|
|
assert torch.equal(cache.project("a", linear, x, compute), linear(x))
|
||
|
|
assert compute.call_count == 2 and not cache._entries
|