acf6620233
Signed-off-by: specture724 <specture724@gmail.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
646 lines
21 KiB
Python
646 lines
21 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
from types import MethodType
|
|
from typing import Any
|
|
from unittest.mock import Mock
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
import torch.nn as nn
|
|
from PIL import Image
|
|
|
|
from vllm_omni.diffusion.data import OmniDiffusionConfig, resolve_model_class_name
|
|
from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata
|
|
from vllm_omni.diffusion.models.magi2.pipeline_magi2 import (
|
|
MAGI2_AUDIO_SAMPLE_RATE,
|
|
MAGI2_MODEL_REVISION,
|
|
Magi2Pipeline,
|
|
_Magi2StagedComponent,
|
|
_resolve_checkpoint_root,
|
|
_validate_native_topology,
|
|
)
|
|
from vllm_omni.diffusion.models.magi2.preview_data_proxy import Magi2DataProxy
|
|
from vllm_omni.diffusion.models.magi2.sampler_magi2 import CFGConfig, Magi2PreviewSampler
|
|
from vllm_omni.diffusion.registry import DiffusionModelRegistry
|
|
from vllm_omni.diffusion.utils.hf_utils import is_diffusion_model
|
|
from vllm_omni.errors import OmniClientError
|
|
from vllm_omni.model_extras.registry import (
|
|
get_extra_body_params,
|
|
should_preserve_reference_image_size,
|
|
)
|
|
|
|
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
|
|
|
|
|
|
@dataclass
|
|
class _SamplingStub:
|
|
width: int | None = None
|
|
height: int | None = None
|
|
fps: float | None = None
|
|
frame_rate: float | None = None
|
|
resolved_frame_rate: float | None = None
|
|
num_frames: int = 1
|
|
num_inference_steps: int | None = None
|
|
num_outputs_per_prompt: int = 1
|
|
seed: int | None = 42
|
|
generator: torch.Generator | list[torch.Generator] | None = None
|
|
extra_args: dict[str, Any] = field(default_factory=dict)
|
|
|
|
|
|
@dataclass
|
|
class _RequestStub:
|
|
prompts: list[Any]
|
|
sampling_params: _SamplingStub
|
|
num_reqs: int = 1
|
|
|
|
|
|
@dataclass
|
|
class _ParallelStub:
|
|
pipeline_parallel_size: int = 1
|
|
data_parallel_size: int = 1
|
|
tensor_parallel_size: int = 1
|
|
sequence_parallel_size: int = 1
|
|
ulysses_degree: int = 1
|
|
ring_degree: int = 1
|
|
allgather_degree: int = 1
|
|
cfg_parallel_size: int = 1
|
|
vae_patch_parallel_size: int = 1
|
|
text_encoder_tp_size: int = 1
|
|
enable_expert_parallel: bool = False
|
|
use_hsdp: bool = False
|
|
|
|
|
|
@dataclass
|
|
class _TopologyStub:
|
|
parallel_config: _ParallelStub = field(default_factory=_ParallelStub)
|
|
diffusion_offload_config: object = None
|
|
enable_cpu_offload: bool = False
|
|
enable_layerwise_offload: bool = False
|
|
enable_distributed_layerwise_offload: bool = False
|
|
dlo_use_allgather: bool = True
|
|
dlo_resident_layers: int = 0
|
|
dlo_host_registration_limit_gib: float = 0.0
|
|
host_weight_runtime_mode: str = "disabled"
|
|
pin_cpu_memory: bool = True
|
|
quantization_config: object = None
|
|
cache_backend: str = "none"
|
|
custom_pipeline_args: dict[str, object] = field(default_factory=dict)
|
|
additional_config: dict[str, object] = field(default_factory=lambda: {"magi2_allow_unsupported_topology": True})
|
|
|
|
|
|
class _FakeNativeRuntime:
|
|
def __init__(self) -> None:
|
|
self.calls: list[dict] = []
|
|
|
|
def evaluate(self, **kwargs):
|
|
self.calls.append(kwargs)
|
|
video = np.zeros((2, 8, 12, 3), dtype=np.uint8)
|
|
audio = np.zeros((32, 2), dtype=np.float32)
|
|
return video, audio
|
|
|
|
|
|
def _pipeline() -> tuple[Magi2Pipeline, _FakeNativeRuntime]:
|
|
pipe = Magi2Pipeline.__new__(Magi2Pipeline)
|
|
nn.Module.__init__(pipe)
|
|
pipe.deterministic = False
|
|
pipe.enable_diffusion_pipeline_profiler = False
|
|
runtime = _FakeNativeRuntime()
|
|
|
|
def evaluate(_self, **kwargs):
|
|
return runtime.evaluate(**kwargs)
|
|
|
|
pipe._evaluate_preview = MethodType(evaluate, pipe)
|
|
return pipe, runtime
|
|
|
|
|
|
def _request(prompt, **sampling_overrides):
|
|
sampling = _SamplingStub()
|
|
for key, value in sampling_overrides.items():
|
|
setattr(sampling, key, value)
|
|
return _RequestStub(prompts=[prompt], sampling_params=sampling)
|
|
|
|
|
|
def _checkpoint_tree(root) -> None:
|
|
files = (
|
|
"preview/model.safetensors.index.json",
|
|
"text_encoder/config.json",
|
|
"text_encoder/model.safetensors.index.json",
|
|
"vae/Wan2.2_VAE.pth",
|
|
"turbo_vae/TurboV3-Wan22-TinyShallow_7_7.json",
|
|
"turbo_vae/checkpoint.ckpt",
|
|
"stable-audio-open-1.0/model_config.json",
|
|
"stable-audio-open-1.0/model.safetensors",
|
|
)
|
|
for relative in files:
|
|
path = root / relative
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.touch()
|
|
|
|
|
|
def test_official_id_detection_and_metadata():
|
|
assert is_diffusion_model("sand-ai/MAGI-2-preview")
|
|
assert is_diffusion_model("https://huggingface.co/sand-ai/MAGI-2-preview")
|
|
assert resolve_model_class_name("sand-ai/MAGI-2-preview") == "Magi2Pipeline"
|
|
assert DiffusionModelRegistry._try_load_model_cls("Magi2Pipeline") is Magi2Pipeline
|
|
metadata = get_diffusion_model_metadata("Magi2Pipeline")
|
|
assert metadata.supports_multimodal_inputs
|
|
assert metadata.max_multimodal_image_inputs == 1
|
|
assert should_preserve_reference_image_size(
|
|
"Magi2Pipeline",
|
|
model="sand-ai/MAGI-2-preview",
|
|
)
|
|
assert {"seconds", "resolution"} <= get_extra_body_params("Magi2Pipeline")
|
|
|
|
|
|
def test_local_preview_signature_detection_without_refiner(tmp_path):
|
|
_checkpoint_tree(tmp_path)
|
|
assert is_diffusion_model(str(tmp_path))
|
|
assert resolve_model_class_name(str(tmp_path)) == "Magi2Pipeline"
|
|
config = OmniDiffusionConfig(model=str(tmp_path))
|
|
config.enrich_config()
|
|
assert config.model_class_name == "Magi2Pipeline"
|
|
assert config.supports_multimodal_inputs
|
|
|
|
|
|
def test_huggingface_url_resolves_pinned_snapshot(tmp_path, monkeypatch):
|
|
_checkpoint_tree(tmp_path)
|
|
api = Mock()
|
|
api.snapshot_download.side_effect = lambda *, repo_id, revision: (
|
|
str(tmp_path)
|
|
if (repo_id, revision) == ("sand-ai/MAGI-2-preview", MAGI2_MODEL_REVISION)
|
|
else pytest.fail("unexpected snapshot request")
|
|
)
|
|
monkeypatch.setattr(
|
|
"vllm_omni.transformers_utils.repo_utils.hf_api",
|
|
lambda: api,
|
|
)
|
|
assert _resolve_checkpoint_root(
|
|
"https://huggingface.co/sand-ai/MAGI-2-preview",
|
|
None,
|
|
) == str(tmp_path.resolve())
|
|
|
|
|
|
def test_forward_uses_native_540p_preview_defaults(monkeypatch):
|
|
pipe, runtime = _pipeline()
|
|
monkeypatch.setattr(
|
|
"vllm_omni.diffusion.models.magi2.pipeline_magi2.current_omni_platform.is_available",
|
|
lambda: False,
|
|
)
|
|
result = pipe(_request("A fox walks through snow"))
|
|
|
|
call = runtime.calls[0]
|
|
assert call["prompt"] == "A fox walks through snow"
|
|
assert call["image"] is None
|
|
assert (call["width"], call["height"]) == (896, 512)
|
|
assert call["num_inference_steps"] == 100
|
|
assert result.output["metadata"]["video"]["fps"] == 12.5
|
|
assert result.output["metadata"]["audio"]["sample_rate"] == MAGI2_AUDIO_SAMPLE_RATE
|
|
assert result.stage_durations["magi2_preview_e2e"] >= 0
|
|
|
|
|
|
def test_forward_maps_272p_i2v_and_output_resize(monkeypatch):
|
|
pipe, runtime = _pipeline()
|
|
monkeypatch.setattr(
|
|
"vllm_omni.diffusion.models.magi2.pipeline_magi2.current_omni_platform.is_available",
|
|
lambda: False,
|
|
)
|
|
image = Image.new("RGB", (16, 9), "white")
|
|
prompt = {
|
|
"prompt": "The first frame begins moving",
|
|
"multi_modal_data": {"image": image},
|
|
}
|
|
result = pipe(
|
|
_request(
|
|
prompt,
|
|
width=448,
|
|
height=256,
|
|
num_inference_steps=1,
|
|
extra_args={
|
|
"output_width": 6,
|
|
"output_height": 4,
|
|
},
|
|
)
|
|
)
|
|
|
|
call = runtime.calls[0]
|
|
assert call["image"] is image
|
|
assert pipe._ensure_figure_token(call["prompt"]) == (
|
|
"The first frame begins moving\nreference_layer:The first frame refers to <Figure 1>"
|
|
)
|
|
assert (call["width"], call["height"]) == (448, 256)
|
|
assert call["num_inference_steps"] == 1
|
|
assert result.output["payload"]["video"].shape == (2, 4, 6, 3)
|
|
|
|
|
|
def test_forward_rejects_conflicting_resolution_and_sampling_geometry(monkeypatch):
|
|
pipe, runtime = _pipeline()
|
|
monkeypatch.setattr(
|
|
"vllm_omni.diffusion.models.magi2.pipeline_magi2.current_omni_platform.is_available",
|
|
lambda: False,
|
|
)
|
|
|
|
with pytest.raises(OmniClientError, match="requires 896x512"):
|
|
pipe(
|
|
_request(
|
|
"A fox walks through snow",
|
|
width=448,
|
|
height=256,
|
|
extra_args={"resolution": "540p"},
|
|
)
|
|
)
|
|
assert not runtime.calls
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("prompt", "extra_args", "message"),
|
|
[
|
|
("", {}, "non-empty prompt"),
|
|
("prompt", {"seconds": 5}, "10-second clips"),
|
|
("prompt", {"duration": 5}, "10-second clips"),
|
|
("prompt", {"resolution": "720p"}, "Unsupported native"),
|
|
("prompt", {"resolution": "1080p"}, "Unsupported native"),
|
|
("prompt", {"resolution": "540p", "use_refiner": True}, "refiner"),
|
|
("prompt", {"output_width": 0, "output_height": 4}, "positive integers"),
|
|
],
|
|
)
|
|
def test_forward_rejects_invalid_preview_requests(
|
|
monkeypatch,
|
|
prompt,
|
|
extra_args,
|
|
message,
|
|
):
|
|
pipe, runtime = _pipeline()
|
|
monkeypatch.setattr(
|
|
"vllm_omni.diffusion.models.magi2.pipeline_magi2.current_omni_platform.is_available",
|
|
lambda: False,
|
|
)
|
|
with pytest.raises(OmniClientError, match=message):
|
|
pipe(_request(prompt, extra_args=extra_args))
|
|
assert not runtime.calls
|
|
|
|
|
|
def test_forward_rejects_zero_inference_steps_before_generation(monkeypatch):
|
|
pipe, runtime = _pipeline()
|
|
monkeypatch.setattr(
|
|
"vllm_omni.diffusion.models.magi2.pipeline_magi2.current_omni_platform.is_available",
|
|
lambda: False,
|
|
)
|
|
with pytest.raises(OmniClientError, match="inference steps must be positive"):
|
|
pipe(_request("prompt", num_inference_steps=0))
|
|
assert not runtime.calls
|
|
|
|
|
|
def test_forward_stops_peak_monitor_when_initial_synchronize_fails(monkeypatch):
|
|
pipe, runtime = _pipeline()
|
|
pipe.enable_diffusion_pipeline_profiler = True
|
|
monitor = Mock(peak_bytes=0)
|
|
monkeypatch.setattr(
|
|
"vllm_omni.diffusion.models.magi2.pipeline_magi2.current_omni_platform.is_cuda",
|
|
lambda: True,
|
|
)
|
|
monkeypatch.setattr(
|
|
"vllm_omni.diffusion.models.magi2.pipeline_magi2.current_omni_platform.is_available",
|
|
lambda: True,
|
|
)
|
|
monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 0)
|
|
monkeypatch.setattr(torch.accelerator, "synchronize", Mock(side_effect=RuntimeError("sync failed")))
|
|
monkeypatch.setattr(
|
|
"vllm_omni.diffusion.models.magi2.pipeline_magi2._PeakReservedMonitor",
|
|
Mock(return_value=monitor),
|
|
)
|
|
|
|
with pytest.raises(RuntimeError, match="sync failed"):
|
|
pipe(_request("A fox walks through snow"))
|
|
|
|
monitor.start.assert_called_once_with()
|
|
monitor.stop.assert_called_once_with()
|
|
assert not runtime.calls
|
|
|
|
|
|
def test_forward_does_not_start_peak_monitor_without_profiler(monkeypatch):
|
|
pipe, runtime = _pipeline()
|
|
monitor_factory = Mock(side_effect=AssertionError("unexpected memory monitor"))
|
|
monkeypatch.setattr(
|
|
"vllm_omni.diffusion.models.magi2.pipeline_magi2.current_omni_platform.is_cuda",
|
|
lambda: True,
|
|
)
|
|
monkeypatch.setattr(
|
|
"vllm_omni.diffusion.models.magi2.pipeline_magi2.current_omni_platform.is_available",
|
|
lambda: True,
|
|
)
|
|
monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 0)
|
|
monkeypatch.setattr(torch.accelerator, "synchronize", Mock())
|
|
monkeypatch.setattr(
|
|
"vllm_omni.diffusion.models.magi2.pipeline_magi2._PeakReservedMonitor",
|
|
monitor_factory,
|
|
)
|
|
|
|
result = pipe(_request("A fox walks through snow"))
|
|
|
|
monitor_factory.assert_not_called()
|
|
assert result.peak_memory_mb == 0.0
|
|
assert len(runtime.calls) == 1
|
|
|
|
|
|
def test_forward_rejects_multiple_images(monkeypatch):
|
|
pipe, _ = _pipeline()
|
|
monkeypatch.setattr(
|
|
"vllm_omni.diffusion.models.magi2.pipeline_magi2.current_omni_platform.is_available",
|
|
lambda: False,
|
|
)
|
|
prompt = {
|
|
"prompt": "animate",
|
|
"multi_modal_data": {"image": [Image.new("RGB", (2, 2)), Image.new("RGB", (2, 2))]},
|
|
}
|
|
with pytest.raises(OmniClientError, match="at most one input image"):
|
|
pipe(_request(prompt))
|
|
|
|
|
|
def test_pipeline_rejects_unknown_initialization_arguments():
|
|
with pytest.raises(TypeError, match=r"Unexpected MAGI-2.*unknown_option"):
|
|
Magi2Pipeline(None, unknown_option=True)
|
|
|
|
|
|
def _topology_config(**overrides):
|
|
parallel = _ParallelStub()
|
|
config = _TopologyStub(parallel_config=parallel)
|
|
for key, value in overrides.items():
|
|
if hasattr(parallel, key):
|
|
setattr(parallel, key, value)
|
|
else:
|
|
setattr(config, key, value)
|
|
return config
|
|
|
|
|
|
def test_native_topology_accepts_qualified_tensor_parallelism():
|
|
_validate_native_topology(_topology_config(tensor_parallel_size=4))
|
|
|
|
|
|
def test_native_topology_rejects_nondivisible_tensor_parallelism():
|
|
config = _topology_config(tensor_parallel_size=3)
|
|
with pytest.raises(ValueError, match="does not divide"):
|
|
_validate_native_topology(config)
|
|
|
|
|
|
def test_native_topology_accepts_single_device_layerwise_with_cpu_staging():
|
|
_validate_native_topology(
|
|
_topology_config(
|
|
enable_layerwise_offload=True,
|
|
additional_config={},
|
|
)
|
|
)
|
|
_validate_native_topology(
|
|
_topology_config(
|
|
enable_cpu_offload=True,
|
|
enable_layerwise_offload=True,
|
|
additional_config={},
|
|
)
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="already stages its auxiliary components"):
|
|
_validate_native_topology(
|
|
_topology_config(
|
|
enable_cpu_offload=True,
|
|
additional_config={},
|
|
)
|
|
)
|
|
|
|
|
|
def test_native_topology_requires_dlo_rank_local_mode():
|
|
with pytest.raises(ValueError, match="dlo-no-use-allgather"):
|
|
_validate_native_topology(_topology_config(enable_distributed_layerwise_offload=True))
|
|
_validate_native_topology(_topology_config(enable_distributed_layerwise_offload=True, dlo_use_allgather=False))
|
|
|
|
|
|
def test_native_topology_accepts_four_device_hsdp_cfg_and_vae_layouts():
|
|
_validate_native_topology(
|
|
_topology_config(
|
|
sequence_parallel_size=4,
|
|
ulysses_degree=4,
|
|
use_hsdp=True,
|
|
)
|
|
)
|
|
_validate_native_topology(
|
|
_topology_config(
|
|
sequence_parallel_size=2,
|
|
ulysses_degree=2,
|
|
cfg_parallel_size=2,
|
|
)
|
|
)
|
|
_validate_native_topology(
|
|
_topology_config(
|
|
sequence_parallel_size=4,
|
|
ulysses_degree=4,
|
|
vae_patch_parallel_size=4,
|
|
)
|
|
)
|
|
|
|
|
|
def test_native_topology_rejects_cfg_data_parallel_combination():
|
|
with pytest.raises(ValueError, match="CFG parallelism.*DLO data parallelism"):
|
|
_validate_native_topology(_topology_config(cfg_parallel_size=2, data_parallel_size=2))
|
|
|
|
|
|
def test_cfg_parallel_branch_adapter_preserves_packed_cfg_math():
|
|
sampler = Magi2PreviewSampler(nn.Identity(), Magi2DataProxy(), device="cpu", dtype=torch.float32)
|
|
latent = torch.arange(8, dtype=torch.float32).reshape(1, 2, 2, 1, 2)
|
|
audio_latent = torch.arange(6, dtype=torch.float32).reshape(1, 3, 2)
|
|
positive_text = torch.ones(1, 3, 4)
|
|
negative_text = torch.full((1, 2, 4), -1.0)
|
|
packed = sampler.prepare_model_input(
|
|
latent=latent,
|
|
audio_latent=audio_latent,
|
|
txt_feat=positive_text,
|
|
null_txt_feat=negative_text,
|
|
t=torch.tensor(500.0),
|
|
cfg_config=CFGConfig(),
|
|
)
|
|
positive, negative = sampler._split_cfg_model_input(packed)
|
|
|
|
assert positive.x_t.shape[0] == negative.x_t.shape[0] == 1
|
|
assert packed.ref_video_feat.shape == (2, 2, 0, 1, 2)
|
|
assert packed.ref_video_feat.numel() == 0
|
|
torch.testing.assert_close(positive.txt_feat[:, :3], positive_text)
|
|
torch.testing.assert_close(negative.txt_feat[:, :2], negative_text)
|
|
|
|
video_pos = torch.full_like(latent, 3.0)
|
|
video_neg = torch.full_like(latent, 1.0)
|
|
audio_pos = torch.full_like(audio_latent, 4.0)
|
|
audio_neg = torch.full_like(audio_latent, 2.0)
|
|
expected = sampler.cfg_velocity(
|
|
(torch.cat((video_pos, video_neg)), torch.cat((audio_pos, audio_neg))),
|
|
5.0,
|
|
7.0,
|
|
CFGConfig(),
|
|
latent,
|
|
audio_latent,
|
|
)
|
|
actual = sampler.combine_cfg_noise(
|
|
(video_pos, audio_pos),
|
|
(video_neg, audio_neg),
|
|
1.0,
|
|
kwargs={
|
|
"video_txt_guidance_scale": 5.0,
|
|
"audio_txt_guidance_scale": 7.0,
|
|
"cfg_config": CFGConfig(),
|
|
"latent": latent,
|
|
"audio_latent": audio_latent,
|
|
},
|
|
)
|
|
torch.testing.assert_close(actual[0], expected[0])
|
|
torch.testing.assert_close(actual[1], expected[1])
|
|
|
|
|
|
def test_decode_audio_preserves_batch_for_released_preview_shape(monkeypatch):
|
|
class ShapeCheckingDecoder(nn.Module):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.seen_shape: tuple[int, ...] | None = None
|
|
|
|
def decode(self, latent: torch.Tensor) -> torch.Tensor:
|
|
self.seen_shape = tuple(latent.shape)
|
|
if self.seen_shape != (1, 64, 250):
|
|
raise AssertionError(f"unexpected decoder input shape {self.seen_shape}")
|
|
return torch.zeros(1, 2, 512_000)
|
|
|
|
decoder = ShapeCheckingDecoder()
|
|
stager = Mock()
|
|
monkeypatch.setattr(
|
|
"vllm_omni.diffusion.models.magi2.pipeline_magi2.PinnedModuleStager",
|
|
Mock(return_value=stager),
|
|
)
|
|
component = _Magi2StagedComponent(
|
|
decoder,
|
|
torch.device("cpu"),
|
|
pin_memory=False,
|
|
)
|
|
pipeline = object.__new__(Magi2Pipeline)
|
|
nn.Module.__init__(pipeline)
|
|
pipeline._is_output_rank = True
|
|
pipeline._offload_aux_after_use = True
|
|
pipeline.device_str = "cpu"
|
|
pipeline.audio_decoder = component
|
|
|
|
def fake_resample(audio: np.ndarray, target_length: int) -> np.ndarray:
|
|
assert audio.shape == (512_000, 2)
|
|
assert target_length == 441_000
|
|
return np.zeros((target_length, 2), dtype=audio.dtype)
|
|
|
|
monkeypatch.setattr("scipy.signal.resample", fake_resample)
|
|
audio = pipeline._decode_audio(torch.zeros(1, 250, 64))
|
|
|
|
assert decoder.seen_shape == (1, 64, 250)
|
|
assert audio is not None
|
|
assert audio.shape == (441_000, 2)
|
|
stager.load.assert_called_once_with()
|
|
stager.offload.assert_called_once_with()
|
|
|
|
|
|
class _TinyMagiTransformer(nn.Module):
|
|
_layerwise_offload_blocks_attrs = ["block"]
|
|
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.block = nn.ModuleList([nn.Linear(2, 2) for _ in range(3)])
|
|
|
|
|
|
class _NoEagerMoveComponent(_Magi2StagedComponent):
|
|
def to(self, *args, **kwargs):
|
|
del args, kwargs
|
|
raise AssertionError("staged MAGI-2 component was moved eagerly")
|
|
|
|
|
|
def _offload_pipeline(monkeypatch):
|
|
stagers = [Mock() for _ in range(4)]
|
|
stager_factory = Mock(side_effect=stagers)
|
|
monkeypatch.setattr(
|
|
"vllm_omni.platforms.current_omni_platform.Stream",
|
|
Mock(return_value=Mock()),
|
|
)
|
|
monkeypatch.setattr(
|
|
"vllm_omni.diffusion.models.magi2.pipeline_magi2.PinnedModuleStager",
|
|
stager_factory,
|
|
)
|
|
pipeline = object.__new__(Magi2Pipeline)
|
|
nn.Module.__init__(pipeline)
|
|
pipeline.transformer = _TinyMagiTransformer()
|
|
components = [
|
|
_NoEagerMoveComponent(
|
|
nn.Linear(2, 2),
|
|
torch.device("cuda:0"),
|
|
pin_memory=False,
|
|
)
|
|
for _ in range(4)
|
|
]
|
|
(
|
|
pipeline.text_encoder,
|
|
pipeline.image_vae,
|
|
pipeline.video_decoder,
|
|
pipeline.audio_decoder,
|
|
) = components
|
|
return pipeline, stagers
|
|
|
|
|
|
class _FakeLayerwiseHook:
|
|
def __init__(self) -> None:
|
|
self._prev_hook = None
|
|
self.current_slot = 0
|
|
self.dp_size = 1
|
|
self.rank_local_mmap = False
|
|
self.prefetch_layer = Mock()
|
|
self.get_weights = Mock()
|
|
|
|
|
|
def test_dlo_no_allgather_enable_preserves_staged_aux_and_streams_blocks(
|
|
monkeypatch,
|
|
):
|
|
from vllm_omni.diffusion.offloader.base import (
|
|
OffloadConfig,
|
|
OffloadStrategy,
|
|
)
|
|
from vllm_omni.diffusion.offloader.distributed_layerwise_backend import (
|
|
DistributedLayerwiseOffloadBackend,
|
|
)
|
|
|
|
pipeline, stagers = _offload_pipeline(monkeypatch)
|
|
hooks: list[_FakeLayerwiseHook] = []
|
|
|
|
def fake_apply(*args, **kwargs):
|
|
del args, kwargs
|
|
hook = _FakeLayerwiseHook()
|
|
hooks.append(hook)
|
|
return hook
|
|
|
|
monkeypatch.setattr(
|
|
"vllm_omni.diffusion.offloader.distributed_layerwise_backend.apply_distributed_block_hook",
|
|
fake_apply,
|
|
)
|
|
backend = DistributedLayerwiseOffloadBackend(
|
|
OffloadConfig(
|
|
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
|
|
pin_cpu_memory=False,
|
|
dlo_use_allgather=False,
|
|
),
|
|
torch.device("cpu"),
|
|
)
|
|
backend._allocate_shared_buffers = Mock(return_value=[{}, {}])
|
|
backend._cleanup_after_loading = Mock()
|
|
backend._release_mmap_handles = Mock()
|
|
|
|
backend.enable(pipeline)
|
|
|
|
assert backend.enabled
|
|
assert backend.dp_size == 1
|
|
assert len(hooks) == len(pipeline.transformer.block)
|
|
assert all(stager.offload.call_count == 1 for stager in stagers)
|
|
assert all(stager.load.call_count == 0 for stager in stagers)
|
|
# First prefetch stays lazy so ranks with different local groups cannot
|
|
# enter collectives in a different order during startup.
|
|
for hook in hooks:
|
|
hook.prefetch_layer.assert_not_called()
|