535bd35b42
Signed-off-by: lishunyang12 <lishunyang12@users.noreply.github.com> Signed-off-by: Armaan Amatya <armaanamatya2014@gmail.com> Co-authored-by: lishunyang12 <lishunyang12@users.noreply.github.com> Co-authored-by: Armaan Amatya <38102001+armaanamatya@users.noreply.github.com>
440 lines
19 KiB
Python
440 lines
19 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from __future__ import annotations
|
|
|
|
import copy
|
|
from dataclasses import asdict
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import torch
|
|
from safetensors.torch import save_file
|
|
|
|
from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
|
|
from vllm_omni.diffusion.models.minimax_h3.adaln_cache import (
|
|
FORMAT_VERSION,
|
|
MODES,
|
|
MiniMaxH3AdalnCache,
|
|
MiniMaxH3RuntimeAdalnCache,
|
|
build_cache,
|
|
canonical_json,
|
|
input_names,
|
|
schedule_contract,
|
|
tensor_digest,
|
|
timestep_plans,
|
|
)
|
|
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline
|
|
|
|
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)
|
|
|
|
|
|
def _contract(mode="t2va", **updates):
|
|
values = dict(
|
|
mode=mode, num_steps=4, base_schedule=(0.999, 0.749, 0.5, 0.25, 0.0), flow_shift=12.0, audio_flow_shift=3.0
|
|
)
|
|
return schedule_contract(**(values | updates))
|
|
|
|
|
|
def _fixture(tmp_path, mode="t2va", adapter=None, *, base_model=False):
|
|
arch = h3.MiniMaxH3DiTArchConfig(
|
|
num_layers=2, hidden_size=8, timestep_input_dim=4, time_embed_hidden_size=8, time_embed_dim=4
|
|
)
|
|
generator = torch.Generator().manual_seed(4)
|
|
weights = {}
|
|
for name, shape in {
|
|
"time_embedder.proj_in.weight": (8, 4),
|
|
"time_embedder.proj_in.bias": (8,),
|
|
"time_embedder.proj_out.weight": (4, 8),
|
|
"time_embedder.proj_out.bias": (4,),
|
|
}.items():
|
|
weights[name] = torch.randn(shape, generator=generator)
|
|
for name in sorted(input_names(arch.num_layers)):
|
|
if name.startswith("time_embedder."):
|
|
continue
|
|
width = (2 if name.startswith("final_layer.") else 18) * arch.hidden_size
|
|
shape = (width, arch.time_embed_dim) if name.endswith("weight") else (width,)
|
|
weights[name] = torch.randn(shape, generator=generator).to(torch.bfloat16)
|
|
variant = "ref2va" if mode.startswith("ref2va") else "fl2va"
|
|
payload, manifest = build_cache(
|
|
arch,
|
|
weights.items(),
|
|
contract=_contract(mode, **({"num_steps": 50, "base_schedule": None} if base_model else {})),
|
|
model_variant=variant,
|
|
adapter_sha256=adapter,
|
|
device=torch.device("cpu"),
|
|
)
|
|
path = tmp_path / "adaln.safetensors"
|
|
_save(path, payload, manifest)
|
|
return arch, weights, payload, manifest, path
|
|
|
|
|
|
def _save(path, payload, manifest):
|
|
save_file(payload, str(path), metadata={"format_version": FORMAT_VERSION, "manifest": canonical_json(manifest)})
|
|
|
|
|
|
def _ready(arch, weights, path, variant="fl2va", adapter=None):
|
|
cache = MiniMaxH3AdalnCache(arch, path=str(path), model_variant=variant)
|
|
cache.bind_adapter(adapter)
|
|
for name, weight in weights.items():
|
|
cache.verify_weight(name, weight)
|
|
cache.finish_loading(torch.device("cpu"))
|
|
return cache
|
|
|
|
|
|
@pytest.mark.parametrize("mode", MODES)
|
|
@pytest.mark.parametrize("base_model", [False, True], ids=["explicit-ladder", "base-h3-50"])
|
|
def test_builder_projection_parity_all_tasks(tmp_path, mode, base_model):
|
|
arch, weights, payload, manifest, path = _fixture(tmp_path, mode, base_model=base_model)
|
|
cache = _ready(arch, weights, path, manifest["model_variant"])
|
|
for i, plan in enumerate(timestep_plans(manifest["schedule"])):
|
|
blocks, final = cache.lookup(plan)
|
|
embedding = payload["time_embeddings"][i, : len(plan)]
|
|
activated = torch.nn.functional.silu(embedding).to(torch.bfloat16)
|
|
for block_index, pieces in enumerate(blocks):
|
|
prefix = f"blocks.{block_index}.adaln_proj.linear"
|
|
expected = torch.nn.functional.linear(activated, weights[prefix + ".weight"], weights[prefix + ".bias"])
|
|
actual = torch.stack(pieces, dim=1).reshape(len(plan), -1)
|
|
assert torch.equal(actual, expected)
|
|
prefix = "final_layer.adaln_proj.linear"
|
|
expected = torch.nn.functional.linear(activated, weights[prefix + ".weight"], weights[prefix + ".bias"])
|
|
assert torch.equal(torch.stack(final, dim=1).reshape(len(plan), -1), expected)
|
|
assert not cache.state_dict()
|
|
|
|
|
|
@pytest.mark.parametrize("change", ["variant", "architecture", "weights", "payload", "plans", "nan", "version"])
|
|
def test_reject_invalid_sidecar(tmp_path, change):
|
|
arch, weights, payload, manifest, path = _fixture(tmp_path)
|
|
payload = {key: value.clone() for key, value in payload.items()}
|
|
manifest = copy.deepcopy(manifest)
|
|
if change == "variant":
|
|
manifest["model_variant"] = "ref2va"
|
|
elif change == "architecture":
|
|
manifest["architecture"]["hidden_size"] += 1
|
|
elif change == "weights":
|
|
manifest["weights"].pop(next(iter(weights)))
|
|
elif change == "payload":
|
|
payload["block_params"][0, 0, 0, 0] += 1
|
|
elif change == "plans":
|
|
payload["plan_timesteps"][0, 0] += 0.01
|
|
manifest["payload"]["plan_timesteps"] = tensor_digest(payload["plan_timesteps"])
|
|
elif change == "nan":
|
|
payload["final_params"][0, 0, 0] = float("nan")
|
|
manifest["payload"]["final_params"] = tensor_digest(payload["final_params"])
|
|
_save(path, payload, manifest)
|
|
if change == "version":
|
|
save_file(payload, str(path), metadata={"format_version": "3"})
|
|
with pytest.raises(ValueError):
|
|
MiniMaxH3AdalnCache(arch, path=str(path), model_variant="fl2va")
|
|
|
|
|
|
def test_weight_adapter_schedule_and_uncovered_timestep_checks(tmp_path):
|
|
arch, weights, _, manifest, path = _fixture(tmp_path, adapter="a" * 64)
|
|
cache = MiniMaxH3AdalnCache(arch, path=str(path), model_variant="fl2va")
|
|
with pytest.raises(ValueError, match="adapter mismatch"):
|
|
cache.bind_adapter(None)
|
|
cache.bind_adapter("a" * 64)
|
|
with pytest.raises(ValueError, match="Missing"):
|
|
cache.finish_loading(torch.device("cpu"))
|
|
name = next(iter(weights))
|
|
with pytest.raises(ValueError, match="effective weight mismatch"):
|
|
cache.verify_weight(name, weights[name] + 1)
|
|
for name, weight in weights.items():
|
|
cache.verify_weight(name, weight)
|
|
cache.finish_loading(torch.device("cpu"))
|
|
for update in (
|
|
{"flow_shift": 10.0},
|
|
{"audio_flow_shift": 4.0},
|
|
{"mode": "fl2va"},
|
|
{"base_schedule": [0.99, 0.74, 0.5, 0.25, 0]},
|
|
):
|
|
with pytest.raises(ValueError):
|
|
cache.check_request(**(manifest["schedule"] | update))
|
|
with pytest.raises(ValueError, match="does not cover"):
|
|
cache.lookup(torch.tensor([0.123]))
|
|
|
|
|
|
def test_adapter_switch_invalidates_runtime_cache(mocker):
|
|
pipeline = object.__new__(MiniMaxH3Pipeline)
|
|
torch.nn.Module.__init__(pipeline)
|
|
cache = mocker.Mock()
|
|
pipeline.transformer = SimpleNamespace(adaln_cache=cache)
|
|
from vllm.lora.request import LoRARequest
|
|
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
|
|
|
sampling = OmniDiffusionSamplingParams(lora_request=LoRARequest("one", 1, "one"), lora_scale=1.0)
|
|
pipeline._prepare_adaln_adapter(sampling)
|
|
pipeline._prepare_adaln_adapter(sampling)
|
|
assert cache.clear.call_count == 1
|
|
sampling.lora_scale = 0.5
|
|
pipeline._prepare_adaln_adapter(sampling)
|
|
sampling.lora_request = None
|
|
pipeline._prepare_adaln_adapter(sampling)
|
|
assert cache.clear.call_count == 3
|
|
|
|
|
|
def test_bad_sidecar_never_drops_runtime_weights(tmp_path):
|
|
arch, weights, _, _, path = _fixture(tmp_path)
|
|
candidate = MiniMaxH3AdalnCache(arch, path=str(path), model_variant="fl2va")
|
|
component = SimpleNamespace(_adaln_sidecar_candidate=candidate)
|
|
changed = list(weights.items())
|
|
changed[0] = changed[0][0], changed[0][1] + 1
|
|
result = list(MiniMaxH3Pipeline._verify_adaln_weights(component, iter(changed)))
|
|
assert len(result) == len(changed)
|
|
assert all(tensor is source for (_, tensor), (_, source) in zip(result, changed))
|
|
assert component._adaln_sidecar_candidate is None
|
|
|
|
|
|
@pytest.mark.parametrize("tp_size", [1, 2])
|
|
def test_native_constructor_enables_cache_by_default(monkeypatch, tp_size):
|
|
from tests.diffusion.models.minimax_h3.test_minimax_h3_quantization import (
|
|
_FakeAttention,
|
|
_FakeLinear,
|
|
_small_od_config,
|
|
)
|
|
|
|
for name in ("ColumnParallelLinear", "RowParallelLinear", "MergedColumnParallelLinear", "QKVParallelLinear"):
|
|
monkeypatch.setattr(h3, name, _FakeLinear)
|
|
monkeypatch.setattr(h3, "Attention", _FakeAttention)
|
|
monkeypatch.setattr(h3, "get_tensor_model_parallel_world_size", lambda: tp_size)
|
|
config = _small_od_config()
|
|
model = h3.MiniMaxH3DiTModel(config, diffusers_weights=False)
|
|
assert model.adaln_cache.max_bytes > 0
|
|
assert model.blocks[0].adaln_proj._adaln_cache is model.adaln_cache
|
|
assert model.final_layer.adaln_proj._adaln_cache is model.adaln_cache
|
|
assert "blocks.0.adaln_proj.linear.weight" in dict(model.named_parameters())
|
|
config.cache_config = {"minimax_h3_adaln_cache": False}
|
|
disabled = h3.MiniMaxH3DiTModel(config, diffusers_weights=False)
|
|
assert disabled.adaln_cache.max_bytes == 0
|
|
|
|
|
|
def test_optional_sidecar_seeds_runtime_projection(tmp_path, mocker):
|
|
arch, weights, payload, _, path = _fixture(tmp_path)
|
|
sidecar = _ready(arch, weights, path)
|
|
runtime = MiniMaxH3RuntimeAdalnCache()
|
|
name = "blocks.0.adaln_proj.linear"
|
|
linear = torch.nn.Linear(arch.time_embed_dim, 18 * arch.hidden_size, dtype=torch.bfloat16)
|
|
with torch.no_grad():
|
|
linear.weight.copy_(weights[name + ".weight"])
|
|
linear.bias.copy_(weights[name + ".bias"])
|
|
runtime.seed(sidecar, {name: linear})
|
|
x = payload["time_embeddings"][1, :2].clone()
|
|
runtime.prepare(x)
|
|
compute = mocker.Mock(side_effect=lambda: linear(torch.nn.functional.silu(x).to(torch.bfloat16)))
|
|
result = runtime.project(name, linear, x, compute)
|
|
assert compute.call_count == 0
|
|
assert torch.equal(result, compute())
|
|
linear.weight.add_(1)
|
|
runtime.project(name, linear, x, compute)
|
|
assert compute.call_count == 2
|
|
|
|
|
|
def test_failed_seed_does_not_attach_sidecar(tmp_path):
|
|
arch, weights, _, _, path = _fixture(tmp_path)
|
|
sidecar = _ready(arch, weights, path)
|
|
runtime = MiniMaxH3RuntimeAdalnCache()
|
|
linear = torch.nn.Linear(arch.time_embed_dim, 18 * arch.hidden_size, dtype=torch.bfloat16)
|
|
with torch.inference_mode():
|
|
unversioned = torch.nn.Linear(arch.time_embed_dim, 2 * arch.hidden_size, dtype=torch.bfloat16)
|
|
|
|
with pytest.raises(RuntimeError, match="Inference tensors do not track version counter"):
|
|
runtime.seed(sidecar, {"blocks.0.adaln_proj.linear": linear, "final_layer.adaln_proj.linear": unversioned})
|
|
|
|
assert runtime.sidecar is None
|
|
assert runtime._sidecar_signatures == {}
|
|
|
|
|
|
@pytest.mark.parametrize("mode", MODES)
|
|
def test_sidecar_prepare_uses_one_fingerprint_without_tensor_scan(tmp_path, mocker, mode):
|
|
import vllm_omni.diffusion.cache.exact_projection_cache as cache_module
|
|
|
|
arch, weights, payload, manifest, path = _fixture(tmp_path, mode=mode, base_model=True)
|
|
sidecar = _ready(arch, weights, path, manifest["model_variant"])
|
|
runtime = MiniMaxH3RuntimeAdalnCache()
|
|
name = "final_layer.adaln_proj.linear"
|
|
linear = torch.nn.Linear(arch.time_embed_dim, 2 * arch.hidden_size, dtype=torch.bfloat16)
|
|
with torch.no_grad():
|
|
linear.weight.copy_(weights[name + ".weight"])
|
|
linear.bias.copy_(weights[name + ".bias"])
|
|
runtime.seed(sidecar, {name: linear})
|
|
count = int(payload["plan_lengths"][-1])
|
|
embedding = payload["time_embeddings"][-1, :count].clone()
|
|
expected = linear(torch.nn.functional.silu(embedding).to(torch.bfloat16))
|
|
comparisons = mocker.spy(torch, "equal")
|
|
fingerprint = mocker.spy(cache_module, "tensor_digest")
|
|
compute = mocker.Mock(side_effect=AssertionError("matching sidecar must seed the projection"))
|
|
|
|
runtime.prepare(embedding)
|
|
actual = runtime.project(name, linear, embedding, compute)
|
|
|
|
comparisons.assert_not_called()
|
|
fingerprint.assert_called_once()
|
|
assert torch.equal(actual, expected)
|
|
compute.assert_not_called()
|
|
|
|
|
|
def test_sidecar_rechecks_numerical_environment_at_projection(tmp_path, mocker):
|
|
arch, weights, payload, _, path = _fixture(tmp_path)
|
|
sidecar = _ready(arch, weights, path)
|
|
runtime = MiniMaxH3RuntimeAdalnCache()
|
|
name = "final_layer.adaln_proj.linear"
|
|
linear = torch.nn.Linear(arch.time_embed_dim, 2 * arch.hidden_size, dtype=torch.bfloat16)
|
|
with torch.no_grad():
|
|
linear.weight.copy_(weights[name + ".weight"])
|
|
linear.bias.copy_(weights[name + ".bias"])
|
|
runtime.seed(sidecar, {name: linear})
|
|
count = int(payload["plan_lengths"][-1])
|
|
embedding = payload["time_embeddings"][-1, :count].clone()
|
|
runtime.prepare(embedding)
|
|
compute = mocker.Mock(side_effect=lambda: linear(torch.nn.functional.silu(embedding).to(torch.bfloat16)))
|
|
precision = torch.get_float32_matmul_precision()
|
|
try:
|
|
torch.set_float32_matmul_precision("high" if precision == "highest" else "highest")
|
|
actual = runtime.project(name, linear, embedding, compute)
|
|
compute.assert_called_once()
|
|
assert torch.equal(actual, compute())
|
|
finally:
|
|
torch.set_float32_matmul_precision(precision)
|
|
|
|
|
|
def test_builder_rejects_missing_or_duplicate_weights(tmp_path):
|
|
arch, weights, _, manifest, _ = _fixture(tmp_path)
|
|
for source in (list(weights.items())[:-1], [*weights.items(), next(iter(weights.items()))]):
|
|
with pytest.raises(ValueError):
|
|
build_cache(
|
|
arch,
|
|
source,
|
|
contract=manifest["schedule"],
|
|
model_variant="fl2va",
|
|
adapter_sha256=None,
|
|
device=torch.device("cpu"),
|
|
)
|
|
|
|
|
|
def test_native_checkpoint_builder_cli(tmp_path, monkeypatch):
|
|
from tools.minimax_h3.build_adaln_cache import main
|
|
|
|
arch, weights, _, _, _ = _fixture(tmp_path)
|
|
root = tmp_path / "transformer"
|
|
root.mkdir()
|
|
(root / "config.json").write_text(canonical_json(asdict(arch)))
|
|
save_file(weights, str(root / "model.safetensors"))
|
|
output = tmp_path / "cli.safetensors"
|
|
argv = [
|
|
"build_adaln_cache",
|
|
"--transformer-path",
|
|
str(root),
|
|
"--output",
|
|
str(output),
|
|
"--model-variant",
|
|
"fl2va",
|
|
"--num-inference-steps",
|
|
"4",
|
|
"--device",
|
|
"cpu",
|
|
]
|
|
monkeypatch.setattr("sys.argv", argv)
|
|
main()
|
|
_ready(arch, weights, output)
|
|
with pytest.raises(ValueError, match="overwrite"):
|
|
main()
|
|
|
|
|
|
@pytest.mark.parametrize("base_model", [False, True], ids=["fasth3-four", "base-h3-50"])
|
|
def test_runtime_projection_matches_sidecar_using_actual_h3_forwards(tmp_path, monkeypatch, mocker, base_model):
|
|
class Linear(torch.nn.Linear):
|
|
def __init__(self, in_features, out_features, *, params_dtype, bias=True, **kwargs):
|
|
super().__init__(in_features, out_features, bias=bias, dtype=params_dtype)
|
|
|
|
def forward(self, x):
|
|
return super().forward(x), None
|
|
|
|
monkeypatch.setattr(h3, "ColumnParallelLinear", Linear)
|
|
monkeypatch.setattr(h3, "RowParallelLinear", Linear)
|
|
arch, weights, payload, manifest, _ = _fixture(tmp_path, base_model=base_model)
|
|
embedder = h3.MiniMaxH3TimeEmbedder(arch, prefix="time_embedder")
|
|
runtime = MiniMaxH3RuntimeAdalnCache()
|
|
proj = h3.MiniMaxH3AdalnProj(
|
|
arch,
|
|
18 * arch.hidden_size,
|
|
None,
|
|
expand_ratio=6,
|
|
modality_num=3,
|
|
prefix="blocks.0.adaln_proj",
|
|
adaln_cache=runtime,
|
|
)
|
|
embedder.load_state_dict(
|
|
{
|
|
name.removeprefix("time_embedder."): value
|
|
for name, value in weights.items()
|
|
if name.startswith("time_embedder.")
|
|
}
|
|
)
|
|
proj.load_state_dict(
|
|
{
|
|
name.removeprefix("blocks.0.adaln_proj."): value
|
|
for name, value in weights.items()
|
|
if name.startswith("blocks.0.adaln_proj.")
|
|
}
|
|
)
|
|
calls = mocker.spy(proj.linear, "forward")
|
|
with torch.no_grad():
|
|
for _ in range(2):
|
|
for i, plan in enumerate(timestep_plans(manifest["schedule"])):
|
|
embedding = embedder(plan)
|
|
assert torch.equal(embedding, payload["time_embeddings"][i, : len(plan)])
|
|
runtime.prepare(embedding)
|
|
actual = torch.stack(proj(embedding), dim=1).reshape(len(plan), -1)
|
|
assert torch.equal(actual, payload["block_params"][i, : len(plan), 0])
|
|
plan_count = len(timestep_plans(manifest["schedule"]))
|
|
assert calls.call_count == plan_count and runtime.hits == plan_count
|
|
|
|
|
|
@pytest.mark.parametrize("enforce_eager", [False, True], ids=["compiled", "eager"])
|
|
@pytest.mark.parametrize(
|
|
"key,mode",
|
|
[("minimax_h3_adaln_cache_path", "t2va"), ("minimax_h3_ref_adaln_cache_path", "ref2va-mixed")],
|
|
)
|
|
def test_optional_sidecar_requires_eager_before_reading_payload(tmp_path, mocker, enforce_eager, key, mode):
|
|
import vllm_omni.diffusion.models.minimax_h3.adaln_cache as cache_module
|
|
|
|
arch, _, _, manifest, path = _fixture(tmp_path, mode=mode)
|
|
open_sidecar = mocker.spy(cache_module, "safe_open")
|
|
pipeline = object.__new__(MiniMaxH3Pipeline)
|
|
torch.nn.Module.__init__(pipeline)
|
|
pipeline.od_config = SimpleNamespace(cache_config={key: str(path)}, enforce_eager=enforce_eager)
|
|
transformer = SimpleNamespace(arch=arch, adaln_cache=MiniMaxH3RuntimeAdalnCache())
|
|
|
|
pipeline._configure_adaln_sidecar(transformer, key, manifest["model_variant"], None, eligible=True)
|
|
|
|
if enforce_eager:
|
|
open_sidecar.assert_called_once()
|
|
assert isinstance(transformer._adaln_sidecar_candidate, MiniMaxH3AdalnCache)
|
|
assert transformer._adaln_sidecar_candidate.path == str(path)
|
|
else:
|
|
open_sidecar.assert_not_called()
|
|
assert not hasattr(transformer, "_adaln_sidecar_candidate")
|
|
# The normal load completion must not install or move a skipped payload.
|
|
pipeline._finish_adaln_sidecar(transformer)
|
|
assert transformer.adaln_cache.sidecar is None
|
|
assert transformer.adaln_cache.max_bytes > 0
|
|
|
|
|
|
def test_corrupt_optional_sidecar_falls_back_without_rejecting_model(tmp_path):
|
|
path = tmp_path / "truncated.safetensors"
|
|
path.write_bytes(b"truncated")
|
|
pipeline = object.__new__(MiniMaxH3Pipeline)
|
|
torch.nn.Module.__init__(pipeline)
|
|
pipeline.od_config = SimpleNamespace(cache_config={"cache": str(path)}, enforce_eager=True)
|
|
transformer = SimpleNamespace(arch=h3.MiniMaxH3DiTArchConfig(), adaln_cache=MiniMaxH3RuntimeAdalnCache())
|
|
pipeline._configure_adaln_sidecar(transformer, "cache", "fl2va", None, eligible=True)
|
|
assert not hasattr(transformer, "_adaln_sidecar_candidate")
|