124 lines
3.9 KiB
Python
124 lines
3.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import pytest
|
|
import torch.nn as nn
|
|
|
|
import vllm_omni.diffusion.compile as compile_module
|
|
from vllm_omni.diffusion.compile import compile_callable_with_fallback, regionally_compile
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
|
|
|
|
|
class _WrappedBlock(nn.Module):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.compile_called = False
|
|
self.forward_compiled = False
|
|
|
|
def compile(self, *args, **kwargs):
|
|
self.compile_called = True
|
|
return self
|
|
|
|
def forward(self, x):
|
|
return x
|
|
|
|
|
|
class _ModelWithWrappedRepeatedBlocks(nn.Module):
|
|
_repeated_blocks = ["OriginalBlock"]
|
|
_layerwise_offload_blocks_attrs = ["transformer_blocks"]
|
|
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.transformer_blocks = nn.ModuleList([_WrappedBlock(), _WrappedBlock()])
|
|
self.other_blocks = nn.ModuleList([_WrappedBlock()])
|
|
|
|
|
|
def test_regionally_compile_matches_wrapped_blocks_by_declared_container_attr(monkeypatch):
|
|
model = _ModelWithWrappedRepeatedBlocks()
|
|
compile_calls = []
|
|
|
|
def _compile(fn, *args, **kwargs):
|
|
compile_calls.append((fn, args, kwargs))
|
|
|
|
def _compiled(*fn_args, **fn_kwargs):
|
|
return f"compiled:{fn(*fn_args, **fn_kwargs)}"
|
|
|
|
return _compiled
|
|
|
|
monkeypatch.setattr(compile_module.torch, "compile", _compile)
|
|
|
|
regionally_compile(model, dynamic=True)
|
|
|
|
assert len(compile_calls) == 2
|
|
assert all(not block.compile_called for block in model.transformer_blocks)
|
|
assert not model.other_blocks[0].compile_called
|
|
assert model.transformer_blocks[0].forward("ok") == "compiled:ok"
|
|
|
|
|
|
def test_regionally_compile_does_not_partially_mutate_on_setup_failure(monkeypatch):
|
|
model = _ModelWithWrappedRepeatedBlocks()
|
|
original_forwards = [block.forward.__func__ for block in model.transformer_blocks]
|
|
compile_calls = 0
|
|
|
|
def _compile(fn, *args, **kwargs):
|
|
nonlocal compile_calls
|
|
compile_calls += 1
|
|
if compile_calls == 2:
|
|
raise RuntimeError("compile setup failed")
|
|
return lambda *fn_args, **fn_kwargs: fn(*fn_args, **fn_kwargs)
|
|
|
|
monkeypatch.setattr(compile_module.torch, "compile", _compile)
|
|
|
|
with pytest.raises(RuntimeError, match="compile setup failed"):
|
|
regionally_compile(model, dynamic=True)
|
|
|
|
assert [block.forward.__func__ for block in model.transformer_blocks] == original_forwards
|
|
|
|
|
|
def test_compile_callable_falls_back_once_after_lazy_failure(monkeypatch, caplog):
|
|
compiled_calls = 0
|
|
eager_calls = 0
|
|
|
|
def eager(value):
|
|
nonlocal eager_calls
|
|
eager_calls += 1
|
|
return value + 1
|
|
|
|
def compile_fn(fn, **kwargs):
|
|
assert fn is eager
|
|
assert kwargs == {"mode": "default"}
|
|
|
|
def compiled(value):
|
|
nonlocal compiled_calls
|
|
compiled_calls += 1
|
|
raise RuntimeError("lazy failure")
|
|
|
|
return compiled
|
|
|
|
monkeypatch.setattr(compile_module.torch, "compile", compile_fn)
|
|
callable_fn, activated = compile_callable_with_fallback(eager, label="test callable", mode="default")
|
|
|
|
assert activated
|
|
assert callable_fn(1) == 2
|
|
assert callable_fn(2) == 3
|
|
assert compiled_calls == 1
|
|
assert eager_calls == 2
|
|
assert "disabling compilation and retrying eagerly" in caplog.text
|
|
|
|
|
|
def test_compile_callable_keeps_eager_on_setup_failure(monkeypatch, caplog):
|
|
def eager(value):
|
|
return value + 1
|
|
|
|
def compile_fn(*args, **kwargs):
|
|
raise RuntimeError("setup failure")
|
|
|
|
monkeypatch.setattr(compile_module.torch, "compile", compile_fn)
|
|
callable_fn, activated = compile_callable_with_fallback(eager, label="test callable")
|
|
|
|
assert not activated
|
|
assert callable_fn is eager
|
|
assert callable_fn(1) == 2
|
|
assert "continuing eagerly" in caplog.text
|