172 lines
6.6 KiB
Python
172 lines
6.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Full-transformer numeric parity for LTX Ulysses sequence parallelism.
|
|
|
|
This FP32/SDPA check verifies that Ulysses preserves the complete transformer
|
|
mathematical path. It does not guarantee pixel-level E2E parity: production
|
|
BF16 and compiled attention kernels can change rounding order, and diffusion
|
|
can amplify that drift across denoising steps.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from tests.helpers.mark import hardware_test
|
|
from vllm_omni.diffusion.config import set_current_diffusion_config
|
|
from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
|
|
from vllm_omni.diffusion.distributed.parallel_state import (
|
|
destroy_distributed_env,
|
|
init_distributed_environment,
|
|
initialize_model_parallel,
|
|
)
|
|
from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig
|
|
from vllm_omni.diffusion.forward_context import get_forward_context, set_forward_context
|
|
from vllm_omni.diffusion.hooks.sequence_parallel import apply_sequence_parallel
|
|
from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel
|
|
from vllm_omni.platforms import current_omni_platform
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.parallel]
|
|
|
|
_FP32_REL_L2 = 1e-5
|
|
_TINY_CONFIG = {
|
|
"in_channels": 4,
|
|
"out_channels": 4,
|
|
"num_attention_heads": 2,
|
|
"attention_head_dim": 4,
|
|
"cross_attention_dim": 8,
|
|
"vae_scale_factors": (1, 1, 1),
|
|
"pos_embed_max_pos": 4,
|
|
"base_height": 2,
|
|
"base_width": 2,
|
|
"audio_in_channels": 4,
|
|
"audio_out_channels": 4,
|
|
"audio_num_attention_heads": 2,
|
|
"audio_attention_head_dim": 4,
|
|
"audio_cross_attention_dim": 8,
|
|
"audio_scale_factor": 1,
|
|
"audio_pos_embed_max_pos": 4,
|
|
"audio_sampling_rate": 16,
|
|
"audio_hop_length": 4,
|
|
"num_layers": 1,
|
|
"caption_channels": 8,
|
|
"rope_type": "split",
|
|
"use_prompt_embeddings": False,
|
|
"use_keyframes_abs_pos_embedding": True,
|
|
}
|
|
|
|
|
|
def _relative_l2(actual: torch.Tensor, expected: torch.Tensor) -> float:
|
|
return ((actual - expected).double().norm() / expected.double().norm()).item()
|
|
|
|
|
|
def _omni_config(sp_size: int) -> OmniDiffusionConfig:
|
|
return OmniDiffusionConfig.from_kwargs(
|
|
model="test",
|
|
dtype=torch.float32,
|
|
parallel_config=DiffusionParallelConfig(
|
|
pipeline_parallel_size=1,
|
|
data_parallel_size=1,
|
|
tensor_parallel_size=1,
|
|
sequence_parallel_size=sp_size,
|
|
ulysses_degree=sp_size,
|
|
ring_degree=1,
|
|
cfg_parallel_size=1,
|
|
),
|
|
diffusion_attention_backend="TORCH_SDPA",
|
|
)
|
|
|
|
|
|
def _inputs(device: torch.device) -> dict[str, torch.Tensor | int | bool]:
|
|
generator = torch.Generator().manual_seed(123)
|
|
return {
|
|
"hidden_states": torch.randn(1, 8, 4, generator=generator, device="cpu").to(device),
|
|
"audio_hidden_states": torch.randn(1, 4, 4, generator=generator, device="cpu").to(device),
|
|
"encoder_hidden_states": torch.randn(1, 3, 8, generator=generator, device="cpu").to(device),
|
|
"audio_encoder_hidden_states": torch.randn(1, 3, 8, generator=generator, device="cpu").to(device),
|
|
"timestep": torch.linspace(0.1, 0.8, 8, device=device).view(1, 8),
|
|
"audio_timestep": torch.linspace(0.2, 0.5, 4, device=device).view(1, 4),
|
|
"keyframes_mask": torch.tensor([[[1.0], [0.0], [1.0], [0.0], [0.0], [1.0], [0.0], [0.0]]], device=device),
|
|
"sigma": torch.tensor([0.1], device=device),
|
|
"audio_sigma": torch.tensor([0.2], device=device),
|
|
"num_frames": 2,
|
|
"height": 2,
|
|
"width": 2,
|
|
"audio_num_frames": 4,
|
|
"return_dict": False,
|
|
}
|
|
|
|
|
|
def _run_transformer_parity(rank: int, world_size: int, master_port: int) -> None:
|
|
device = torch.device(f"{current_omni_platform.device_type}:{rank}")
|
|
current_omni_platform.set_device(device)
|
|
os.environ.update(
|
|
{
|
|
"RANK": str(rank),
|
|
"LOCAL_RANK": str(rank),
|
|
"WORLD_SIZE": str(world_size),
|
|
"MASTER_ADDR": "localhost",
|
|
"MASTER_PORT": str(master_port),
|
|
}
|
|
)
|
|
|
|
try:
|
|
init_distributed_environment()
|
|
initialize_model_parallel(sequence_parallel_size=world_size, ulysses_degree=world_size)
|
|
|
|
sp_config = _omni_config(world_size)
|
|
torch.manual_seed(42)
|
|
with set_forward_context(omni_diffusion_config=sp_config), set_current_diffusion_config(sp_config):
|
|
model = LTX2VideoTransformer3DModel(**_TINY_CONFIG).to(device).eval()
|
|
with torch.no_grad():
|
|
for _, parameter in sorted(model.named_parameters()):
|
|
torch.nn.init.normal_(parameter, mean=0.0, std=0.02)
|
|
|
|
inputs = _inputs(device)
|
|
reference_config = _omni_config(1)
|
|
with torch.no_grad(), set_forward_context(omni_diffusion_config=reference_config):
|
|
expected_video, expected_audio = model(**inputs)
|
|
|
|
apply_sequence_parallel(
|
|
model,
|
|
SequenceParallelConfig(ulysses_degree=world_size, ring_degree=1),
|
|
model._sp_plan,
|
|
)
|
|
with torch.no_grad(), set_forward_context(omni_diffusion_config=sp_config):
|
|
get_forward_context().sp_plan_hooks_applied = True
|
|
actual_video, actual_audio = model(**inputs)
|
|
|
|
v2a_attention = model.transformer_blocks[0].video_to_audio_attn.attn
|
|
assert v2a_attention._video_to_audio_strategy is not None
|
|
assert torch.isfinite(expected_video).all()
|
|
assert torch.isfinite(expected_audio).all()
|
|
for modality, actual, expected in (
|
|
("video", actual_video, expected_video),
|
|
("audio", actual_audio, expected_audio),
|
|
):
|
|
drift = _relative_l2(actual, expected)
|
|
assert drift <= _FP32_REL_L2, (
|
|
f"LTX SP{world_size} {modality} output drifted rel_l2={drift:.3e} from SP1 (bound {_FP32_REL_L2:.1e})"
|
|
)
|
|
finally:
|
|
destroy_distributed_env()
|
|
|
|
|
|
@pytest.mark.full_model
|
|
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=2)
|
|
def test_ltx_ulysses_transformer_matches_sp1(unused_tcp_port) -> None:
|
|
"""Exercise real plan hooks, RoPE/timestep sharding, and attention routing."""
|
|
world_size = 2
|
|
device_count = current_omni_platform.device_count()
|
|
assert device_count >= world_size, (
|
|
f"LTX Ulysses parity requires {world_size} accelerator devices; only {device_count} are visible"
|
|
)
|
|
torch.multiprocessing.spawn(
|
|
_run_transformer_parity,
|
|
args=(world_size, unused_tcp_port),
|
|
nprocs=world_size,
|
|
)
|