275 lines
13 KiB
Python
275 lines
13 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from dataclasses import dataclass
|
|
from datetime import timedelta
|
|
|
|
import pytest
|
|
import torch
|
|
import torch.distributed as dist
|
|
import torch.multiprocessing as mp
|
|
from diffusers import AutoencoderKL, AutoencoderKLFlux2
|
|
from torch import nn
|
|
|
|
from vllm_omni.config.omni_config import OmniStageDiffusionParallelConfig
|
|
from vllm_omni.diffusion import registry
|
|
from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
|
|
from vllm_omni.diffusion.distributed.autoencoders import autoencoder_kl as kl_module
|
|
from vllm_omni.diffusion.distributed.autoencoders import distributed_vae_executor as executor_module
|
|
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl import DistributedAutoencoderKL
|
|
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_flux2 import DistributedAutoencoderKLFlux2
|
|
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan
|
|
from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import DistributedVaeExecutor
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion, pytest.mark.parallel]
|
|
|
|
|
|
class _RecordingDecoder(nn.Module):
|
|
def __init__(self, dtype):
|
|
super().__init__()
|
|
self.conv = nn.Conv2d(1, 3, kernel_size=1, bias=False, dtype=dtype)
|
|
with torch.no_grad():
|
|
self.conv.weight.copy_(torch.arange(1, 4, dtype=dtype).reshape(3, 1, 1, 1))
|
|
self.batch_sizes = []
|
|
|
|
def forward(self, z):
|
|
self.batch_sizes.append(z.shape[0])
|
|
return self.conv(z)
|
|
|
|
|
|
def _make_vae(vae_type=DistributedAutoencoderKLFlux2, dtype=torch.float32):
|
|
"""Keep the native decode methods; replace only the learned decoder."""
|
|
vae = vae_type.__new__(vae_type)
|
|
nn.Module.__init__(vae)
|
|
vae.register_to_config(use_post_quant_conv=False)
|
|
vae.decoder = _RecordingDecoder(dtype)
|
|
vae.post_quant_conv = None
|
|
vae.use_slicing = False
|
|
vae.use_tiling = False
|
|
return vae
|
|
|
|
|
|
@dataclass
|
|
class _WorldGroup:
|
|
device_group: dist.ProcessGroup
|
|
cpu_group: dist.ProcessGroup
|
|
|
|
|
|
def _batch_worker(rank, rendezvous):
|
|
torch.set_num_threads(1)
|
|
dist.init_process_group("gloo", init_method=rendezvous, rank=rank, world_size=4, timeout=timedelta(seconds=60))
|
|
try:
|
|
with pytest.MonkeyPatch.context() as patch:
|
|
patch.setattr(executor_module, "get_world_group", lambda: _WorldGroup(dist.group.WORLD, dist.group.WORLD))
|
|
for name in (
|
|
"get_data_parallel_world_size",
|
|
"get_pipeline_parallel_world_size",
|
|
"get_classifier_free_guidance_world_size",
|
|
):
|
|
patch.setattr(kl_module, name, lambda: 1)
|
|
for vae_type, native_type in (
|
|
(DistributedAutoencoderKL, AutoencoderKL),
|
|
(DistributedAutoencoderKLFlux2, AutoencoderKLFlux2),
|
|
):
|
|
for dtype, autocast in ((torch.float32, False), (torch.bfloat16, False), (torch.float32, True)):
|
|
vae = _make_vae(vae_type, dtype)
|
|
vae.init_distributed()
|
|
vae.set_parallel_size(4, mode="batch")
|
|
vae.use_tiling = True
|
|
assert not vae.is_distributed_enabled() # Spatial encode/decode stays disabled.
|
|
vae.use_tiling = False
|
|
# Reuse the same instance with changing batches and contents.
|
|
for batch, return_dict, expected_sizes in (
|
|
(2, False, (1, 1, 0, 0)),
|
|
(5, True, (2, 1, 1, 1)),
|
|
(1, True, (1, 1, 1, 1)),
|
|
):
|
|
z = torch.arange(batch * 6, dtype=dtype).reshape(batch, 1, 2, 3) + batch
|
|
with torch.autocast("cpu", dtype=torch.bfloat16, enabled=autocast):
|
|
expected = native_type.decode(vae, z, return_dict=False)[0]
|
|
vae.decoder.batch_sizes.clear()
|
|
if vae_type is DistributedAutoencoderKL and not return_dict:
|
|
result = vae.decode(z, False, torch.Generator().manual_seed(0))
|
|
else:
|
|
result = vae.decode(z, return_dict=return_dict)
|
|
actual = result.sample if return_dict else result[0]
|
|
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
|
|
assert actual.dtype == expected.dtype
|
|
assert sum(vae.decoder.batch_sizes) == expected_sizes[rank]
|
|
assert not vae.use_tiling
|
|
|
|
# Degree one bypasses collectives and keeps native slicing.
|
|
vae.set_parallel_size(1, mode="batch")
|
|
vae.use_slicing = True
|
|
z = torch.ones(2, 1, 2, 3, dtype=dtype)
|
|
vae.decoder.batch_sizes.clear()
|
|
if vae_type is DistributedAutoencoderKL:
|
|
actual = vae.decode(z, False, torch.Generator().manual_seed(0))[0]
|
|
else:
|
|
actual = vae.decode(z, return_dict=False)[0]
|
|
assert vae.decoder.batch_sizes == [1, 1]
|
|
expected = native_type.decode(vae, z, return_dict=False)[0]
|
|
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
|
|
|
|
# Execute native tiled decode, including the B1 fallback and
|
|
# uneven batches with idle ranks; preserving the flag alone
|
|
# does not prove that tiling is still used.
|
|
vae = _make_vae(vae_type)
|
|
vae.init_distributed()
|
|
vae.set_parallel_size(4, mode="batch")
|
|
vae.use_tiling = True
|
|
vae.tile_latent_min_size = vae.tile_sample_min_size = 4
|
|
vae.tile_overlap_factor = 0.5
|
|
patch.setattr(
|
|
vae, "_strategy_select", lambda _: pytest.fail("Batch mode must not distribute spatial tiles")
|
|
)
|
|
for batch, expected_sizes in (
|
|
(1, (1, 1, 1, 1)),
|
|
(2, (1, 1, 0, 0)),
|
|
(5, (2, 1, 1, 1)),
|
|
):
|
|
z = torch.arange(batch * 36, dtype=torch.float32).reshape(batch, 1, 6, 6)
|
|
expected = native_type.decode(vae, z, return_dict=False)[0]
|
|
vae.decoder.batch_sizes.clear()
|
|
actual = vae.decode(z).sample
|
|
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
|
|
assigned_batch = expected_sizes[rank]
|
|
if assigned_batch:
|
|
# More than one decoder call proves native tiling ran.
|
|
assert len(vae.decoder.batch_sizes) > 1
|
|
assert set(vae.decoder.batch_sizes) == {assigned_batch}
|
|
else:
|
|
assert vae.decoder.batch_sizes == []
|
|
assert vae.use_tiling
|
|
finally:
|
|
dist.destroy_process_group()
|
|
|
|
|
|
@pytest.mark.skipif(not dist.is_gloo_available(), reason="Requires the CPU Gloo backend")
|
|
def test_batch_decode_matches_native_on_all_ranks(tmp_path):
|
|
mp.spawn(_batch_worker, args=(f"file://{tmp_path / 'rendezvous'}",), nprocs=4, join=True)
|
|
|
|
|
|
def test_batch_decode_without_distributed_initialization_uses_native(monkeypatch, mocker):
|
|
vae = _make_vae()
|
|
executor = mocker.Mock(spec=DistributedVaeExecutor)
|
|
executor.parallel_mode = "batch"
|
|
executor.parallel_size = executor.world_size = 4
|
|
vae.distributed_executor = executor
|
|
monkeypatch.setattr(dist, "is_initialized", lambda: False)
|
|
z = torch.ones(2, 1, 2, 3)
|
|
expected = AutoencoderKLFlux2.decode(vae, z, return_dict=False)[0]
|
|
torch.testing.assert_close(vae.decode(z).sample, expected, rtol=0, atol=0)
|
|
executor.execute.assert_not_called()
|
|
|
|
|
|
def test_flux2_batch_encode_keeps_native_tiling(monkeypatch, mocker):
|
|
vae = _make_vae()
|
|
vae.register_to_config(use_quant_conv=False)
|
|
vae.encoder = _RecordingDecoder(torch.float32)
|
|
vae.quant_conv = None
|
|
vae.use_tiling = True
|
|
vae.tile_sample_min_size = vae.tile_latent_min_size = 4
|
|
vae.tile_overlap_factor = 0.5
|
|
executor = mocker.Mock(spec=DistributedVaeExecutor)
|
|
executor.parallel_mode = "batch"
|
|
executor.parallel_size = executor.world_size = 4
|
|
executor.group = None
|
|
vae.distributed_executor = executor
|
|
monkeypatch.setattr(dist, "is_initialized", lambda: True)
|
|
monkeypatch.setattr(dist, "get_world_size", lambda group: 4)
|
|
x = torch.arange(36, dtype=torch.float32).reshape(1, 1, 6, 6)
|
|
expected = AutoencoderKLFlux2._encode(vae, x)
|
|
native_tiling = mocker.spy(vae, "_tiled_encode")
|
|
actual = vae._encode(x)
|
|
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
|
|
native_tiling.assert_called_once_with(x)
|
|
executor.execute.assert_not_called()
|
|
|
|
|
|
@pytest.mark.parametrize("config_type", [DiffusionParallelConfig, OmniStageDiffusionParallelConfig])
|
|
@pytest.mark.parametrize("axis", ["data_parallel_size", "pipeline_parallel_size", "cfg_parallel_size"])
|
|
def test_batch_mode_rejects_separate_request_or_pipeline_groups(config_type, axis):
|
|
with pytest.raises(ValueError, match="batch.*DP, PP, and CFG"):
|
|
config_type(vae_parallel_mode="batch", **{axis: 2})
|
|
|
|
|
|
@pytest.mark.parametrize("config_type", [DiffusionParallelConfig, OmniStageDiffusionParallelConfig])
|
|
def test_batch_mode_accepts_tensor_parallel_group(config_type):
|
|
config = config_type(tensor_parallel_size=4, vae_patch_parallel_size=4, vae_parallel_mode="batch")
|
|
assert config.world_size == 4
|
|
|
|
|
|
def test_batch_mode_rejects_implicitly_resolved_data_parallelism():
|
|
config = DiffusionParallelConfig(vae_parallel_mode="batch", data_parallel_size=None)
|
|
with pytest.raises(ValueError, match="batch.*DP"):
|
|
config.resolve_data_parallel_size(2)
|
|
|
|
|
|
@pytest.mark.parametrize("tiling", [False, True])
|
|
def test_registry_preserves_requested_tiling_in_batch_mode(monkeypatch, mocker, tiling):
|
|
vae = _make_vae()
|
|
vae.distributed_executor = mocker.Mock(spec=DistributedVaeExecutor)
|
|
|
|
class Pipeline(nn.Module):
|
|
def __init__(self, od_config):
|
|
super().__init__()
|
|
self.vae = vae
|
|
|
|
monkeypatch.setattr(registry.DiffusionModelRegistry, "_try_load_model_cls", lambda _: Pipeline)
|
|
config = OmniDiffusionConfig(
|
|
model_class_name="test_batch_vae",
|
|
parallel_config=DiffusionParallelConfig(
|
|
tensor_parallel_size=4, vae_patch_parallel_size=4, vae_parallel_mode="batch"
|
|
),
|
|
vae_use_tiling=tiling,
|
|
vae_use_slicing=True,
|
|
vae_fast_path="off",
|
|
)
|
|
model = registry.initialize_model(config)
|
|
assert model.vae.use_tiling is tiling
|
|
assert model.vae.use_slicing is True
|
|
assert config.vae_use_tiling is tiling
|
|
vae.distributed_executor.set_parallel_size.assert_called_once_with(4, mode="batch")
|
|
|
|
|
|
@pytest.mark.parametrize("vae_type", [nn.Identity, DistributedAutoencoderKLWan])
|
|
def test_registry_rejects_unsupported_batch_vae(monkeypatch, vae_type):
|
|
class Pipeline(nn.Module):
|
|
def __init__(self, od_config):
|
|
super().__init__()
|
|
self.vae = vae_type.__new__(vae_type)
|
|
nn.Module.__init__(self.vae)
|
|
|
|
monkeypatch.setattr(registry.DiffusionModelRegistry, "_try_load_model_cls", lambda _: Pipeline)
|
|
config = OmniDiffusionConfig(
|
|
model_class_name="unsupported_batch_vae",
|
|
parallel_config=DiffusionParallelConfig(vae_parallel_mode="batch"),
|
|
)
|
|
with pytest.raises(ValueError, match="batch.*AutoencoderKL"):
|
|
registry.initialize_model(config)
|
|
|
|
|
|
@pytest.mark.parametrize("mode", ["tile", "batch"])
|
|
def test_klein_uses_distributed_vae_only_for_batch_mode(monkeypatch, tmp_path, mode):
|
|
from vllm_omni.diffusion.models.flux2_klein import pipeline_flux2_klein as module
|
|
|
|
class StopAtVaeLoadError(Exception):
|
|
pass
|
|
|
|
def load_component(factory, *args, subfolder, **kwargs):
|
|
if subfolder == "text_encoder":
|
|
return nn.Identity()
|
|
expected = DistributedAutoencoderKLFlux2 if mode == "batch" else AutoencoderKLFlux2
|
|
assert factory.__self__ is expected
|
|
raise StopAtVaeLoadError
|
|
|
|
monkeypatch.setattr(module, "get_local_device", lambda: torch.device("cpu"))
|
|
monkeypatch.setattr(module, "prefetch_subfolders", lambda *args, **kwargs: None)
|
|
monkeypatch.setattr(module, "from_pretrained_with_prefetch", load_component)
|
|
monkeypatch.setattr(module.FlowMatchEulerDiscreteScheduler, "from_pretrained", lambda *args, **kwargs: None)
|
|
monkeypatch.setattr(module.Qwen2TokenizerFast, "from_pretrained", lambda *args, **kwargs: None)
|
|
config = OmniDiffusionConfig(model=str(tmp_path), parallel_config=DiffusionParallelConfig(vae_parallel_mode=mode))
|
|
with pytest.raises(StopAtVaeLoadError):
|
|
module.Flux2KleinPipeline(od_config=config)
|