78 lines
3.3 KiB
Python
78 lines
3.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from tests.diffusion.models.magi2.test_native_packing import (
|
|
_longer_text_tensors,
|
|
_sampler_tensors,
|
|
_tiny_model,
|
|
_tiny_sampler,
|
|
)
|
|
from tests.helpers.mark import hardware_marks
|
|
from vllm_omni.diffusion.models.magi2.sampler_magi2 import CFGConfig
|
|
|
|
pytestmark = [
|
|
pytest.mark.diffusion,
|
|
pytest.mark.core_model,
|
|
*hardware_marks(res={"cuda": "L4"}, num_cards=1),
|
|
pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA"),
|
|
]
|
|
|
|
|
|
def _cuda_steps(sampler, tensors=_sampler_tensors):
|
|
timesteps = (torch.tensor([900.0]), torch.tensor([600.0]), torch.tensor([300.0]))
|
|
return [
|
|
sampler.prepare_model_input(
|
|
**{name: value.cuda() for name, value in tensors(seed).items()},
|
|
t=t,
|
|
cfg_config=CFGConfig(),
|
|
)
|
|
for seed, t in enumerate(timesteps)
|
|
]
|
|
|
|
|
|
def test_inductor_regions_match_eager_in_deterministic_mode(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("MAGI2_DETERMINISTIC", "1")
|
|
model = _tiny_model(params_dtype=torch.bfloat16).cuda()
|
|
sampler = _tiny_sampler(model)
|
|
steps = _cuda_steps(sampler)
|
|
expected = [sampler.forward(step) for step in steps]
|
|
new_request_step = _cuda_steps(sampler)[0]
|
|
longer_text_step = _cuda_steps(sampler, _longer_text_tensors)[0]
|
|
expected_new_request = sampler.forward(new_request_step)
|
|
expected_longer_text = sampler.forward(longer_text_step)
|
|
|
|
torch._dynamo.reset()
|
|
graphs_before = torch._dynamo.utils.counters["stats"]["unique_graphs"]
|
|
# The pipeline compiles with inductor emulating eager's intermediate
|
|
# low-precision casts, so fused bf16 chains keep eager's rounding.
|
|
model.compile_regions(fullgraph=True, dynamic=True, options={"emulate_precision_casts": True})
|
|
first = sampler.forward(steps[0])
|
|
with torch._dynamo.config.patch(error_on_recompile=True):
|
|
compiled = [first, sampler.forward(steps[1]), sampler.forward(steps[2])]
|
|
repeated = [sampler.forward(step) for step in steps]
|
|
# A second request repacks the same shapes; a third changes the
|
|
# text length and therefore every token count.
|
|
compiled_new_request = sampler.forward(new_request_step)
|
|
compiled_longer_text = sampler.forward(longer_text_step)
|
|
|
|
# Emulated precision casts restore eager's rounding boundaries; what is
|
|
# left is one-ulp transcendental noise that some Triton backends carry.
|
|
for (video, audio), (video_ref, audio_ref) in zip(compiled, expected, strict=True):
|
|
torch.testing.assert_close(video, video_ref, atol=1e-6, rtol=0)
|
|
torch.testing.assert_close(audio, audio_ref, atol=1e-6, rtol=0)
|
|
for (video, audio), (video_again, audio_again) in zip(compiled, repeated, strict=True):
|
|
assert torch.equal(video, video_again)
|
|
assert torch.equal(audio, audio_again)
|
|
for (video, audio), (video_ref, audio_ref) in (
|
|
(compiled_new_request, expected_new_request),
|
|
(compiled_longer_text, expected_longer_text),
|
|
):
|
|
torch.testing.assert_close(video, video_ref, atol=1e-6, rtol=0)
|
|
torch.testing.assert_close(audio, audio_ref, atol=1e-6, rtol=0)
|
|
assert torch._dynamo.utils.counters["stats"]["unique_graphs"] - graphs_before == 7
|