240 lines
10 KiB
Python
240 lines
10 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import weakref
|
|
|
|
import pytest
|
|
import torch
|
|
from torch import nn
|
|
|
|
from vllm_omni.diffusion.distributed.sp_plan import (
|
|
SequenceParallelConfig,
|
|
SequenceParallelInput,
|
|
SequenceParallelPartialInput,
|
|
)
|
|
from vllm_omni.diffusion.forward_context import get_forward_context, set_forward_context
|
|
from vllm_omni.diffusion.hooks.base import HookRegistry
|
|
from vllm_omni.diffusion.hooks.sequence_parallel import SequenceParallelSplitHook, apply_sequence_parallel
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion, pytest.mark.parallel]
|
|
|
|
|
|
@pytest.fixture
|
|
def parallel_state(monkeypatch):
|
|
from vllm_omni.diffusion.attention import selector
|
|
from vllm_omni.diffusion.distributed import parallel_state, sp_sharding
|
|
|
|
class MaskBackend:
|
|
@staticmethod
|
|
def supports_attention_mask(spec):
|
|
return True
|
|
|
|
def configure(size=4, rank=1):
|
|
monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: size)
|
|
monkeypatch.setattr(parallel_state, "get_sequence_parallel_rank", lambda: rank)
|
|
monkeypatch.setattr(parallel_state, "get_ring_parallel_world_size", lambda: 1)
|
|
monkeypatch.setattr(sp_sharding, "get_sequence_parallel_world_size", lambda: size)
|
|
monkeypatch.setattr(sp_sharding, "get_sequence_parallel_rank", lambda: rank)
|
|
monkeypatch.setattr(selector, "get_attn_backend_for_capability", lambda **kwargs: MaskBackend())
|
|
|
|
configure()
|
|
return configure
|
|
|
|
|
|
@pytest.mark.parametrize("sequence_length", [16, 17])
|
|
@pytest.mark.parametrize("batch", [1, 2])
|
|
@pytest.mark.parametrize("rank", [0, 3])
|
|
@torch.inference_mode()
|
|
def test_cosmos3_real_sp_hook_releases_full_embedding(parallel_state, sequence_length, batch, rank):
|
|
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3GenSPPrepare, Cosmos3VFMTransformer
|
|
|
|
parallel_state(rank=rank)
|
|
model = nn.Module()
|
|
model.gen_sp_prepare = Cosmos3GenSPPrepare()
|
|
metadata = Cosmos3VFMTransformer._sp_plan["gen_sp_prepare"]
|
|
assert metadata[0].clone_shard
|
|
assert not metadata[1].clone_shard and not metadata[2].clone_shard
|
|
apply_sequence_parallel(model, SequenceParallelConfig(ulysses_degree=4), {"gen_sp_prepare": metadata})
|
|
|
|
def prepare():
|
|
hidden = torch.arange(batch * sequence_length * 8, dtype=torch.bfloat16).reshape(batch, sequence_length, 8)
|
|
full_ref = weakref.ref(hidden)
|
|
full_pointer = hidden.untyped_storage().data_ptr()
|
|
cos = torch.ones(batch, sequence_length, 1, 2)
|
|
sin = torch.zeros_like(cos)
|
|
return model.gen_sp_prepare(hidden, cos, sin), full_ref, full_pointer
|
|
|
|
for _ in range(2):
|
|
with set_forward_context():
|
|
(hidden, cos, sin), full_ref, full_pointer = prepare()
|
|
# Check the real backing allocation, not only Tensor.data_ptr():
|
|
# nonzero-offset and inference-mode views can hide the shared base.
|
|
assert full_ref() is None
|
|
assert hidden.untyped_storage().data_ptr() != full_pointer
|
|
assert hidden.untyped_storage().nbytes() == hidden.numel() * hidden.element_size()
|
|
assert hidden.is_contiguous()
|
|
expected = torch.arange(batch * sequence_length * 8, dtype=torch.bfloat16).reshape(
|
|
batch, sequence_length, 8
|
|
)
|
|
padding = -sequence_length % 4
|
|
expected = torch.nn.functional.pad(expected, (0, 0, 0, padding)).chunk(4, dim=1)[rank]
|
|
torch.testing.assert_close(hidden, expected, rtol=0, atol=0)
|
|
assert cos.shape[1] == sin.shape[1] == hidden.shape[1]
|
|
|
|
# Once the next layer replaces the shard, the hook must not keep
|
|
# even the smaller initial activation alive through later layers.
|
|
shard_ref = weakref.ref(hidden)
|
|
hidden = hidden + 1
|
|
assert shard_ref() is None
|
|
|
|
|
|
@pytest.mark.parametrize("clone_shard", [False, True])
|
|
@pytest.mark.parametrize("world_size", [1, 4])
|
|
def test_shard_storage_policy_is_opt_in_and_skips_single_rank(parallel_state, clone_shard, world_size):
|
|
parallel_state(size=world_size, rank=0)
|
|
hook = SequenceParallelSplitHook({}, SequenceParallelConfig(ulysses_degree=4))
|
|
tensor = torch.arange(128).reshape(1, 16, 8)
|
|
spec = SequenceParallelInput(split_dim=1, clone_shard=clone_shard)
|
|
with set_forward_context():
|
|
shard = hook._prepare_sp_input(tensor, spec)
|
|
torch.testing.assert_close(shard, tensor.chunk(world_size, dim=1)[0])
|
|
owns_storage = shard.untyped_storage().data_ptr() != tensor.untyped_storage().data_ptr()
|
|
assert owns_storage == (clone_shard and world_size > 1)
|
|
if world_size == 1:
|
|
assert shard is tensor
|
|
|
|
|
|
def test_owned_shard_preserves_gradients(parallel_state):
|
|
parallel_state(rank=1)
|
|
hook = SequenceParallelSplitHook({}, SequenceParallelConfig(ulysses_degree=4))
|
|
tensor = torch.randn(2, 16, 8, requires_grad=True)
|
|
with set_forward_context():
|
|
shard = hook._prepare_sp_input(tensor, SequenceParallelInput(split_dim=1, clone_shard=True))
|
|
assert shard.untyped_storage().data_ptr() != tensor.untyped_storage().data_ptr()
|
|
assert shard.untyped_storage().nbytes() == shard.numel() * shard.element_size()
|
|
shard.sum().backward()
|
|
expected = torch.zeros_like(tensor)
|
|
expected[:, 4:8] = 1
|
|
torch.testing.assert_close(tensor.grad, expected, rtol=0, atol=0)
|
|
|
|
|
|
@pytest.mark.parametrize("sequence_length", [16, 17])
|
|
@pytest.mark.parametrize("batch", [1, 2])
|
|
@pytest.mark.parametrize("keyword", [False, True])
|
|
@pytest.mark.parametrize("shard_group", [None, "video"])
|
|
@torch.inference_mode()
|
|
def test_pre_forward_clone_owns_storage_and_preserves_shard_metadata(
|
|
parallel_state, sequence_length, batch, keyword, shard_group
|
|
):
|
|
parallel_state(rank=1)
|
|
full_refs = []
|
|
full_pointers = []
|
|
|
|
class Module(nn.Module):
|
|
def forward(self, x):
|
|
assert x.untyped_storage().data_ptr() != full_pointers[-1]
|
|
assert x.untyped_storage().nbytes() == x.numel() * x.element_size()
|
|
assert x.is_contiguous()
|
|
return x
|
|
|
|
module = Module()
|
|
spec = SequenceParallelInput(
|
|
split_dim=1,
|
|
split_output=False,
|
|
auto_pad=True,
|
|
clone_shard=True,
|
|
shard_group=shard_group,
|
|
)
|
|
apply_sequence_parallel(module, SequenceParallelConfig(ulysses_degree=4), {"": {"x": spec}})
|
|
|
|
def run():
|
|
tensor = torch.arange(batch * sequence_length * 8, dtype=torch.bfloat16).reshape(batch, sequence_length, 8)
|
|
full_refs.append(weakref.ref(tensor))
|
|
full_pointers.append(tensor.untyped_storage().data_ptr())
|
|
return module(x=tensor) if keyword else module(tensor)
|
|
|
|
for _ in range(2):
|
|
with set_forward_context():
|
|
ctx = get_forward_context()
|
|
if shard_group is not None:
|
|
ctx.sp_shard_metadata[shard_group] = 999 # Stale metadata must be replaced.
|
|
shard = run()
|
|
assert all(ref() is None for ref in full_refs)
|
|
expected = torch.arange(batch * sequence_length * 8, dtype=torch.bfloat16).reshape(
|
|
batch, sequence_length, 8
|
|
)
|
|
padding = -sequence_length % 4
|
|
expected = torch.nn.functional.pad(expected, (0, 0, 0, padding)).chunk(4, dim=1)[1]
|
|
torch.testing.assert_close(shard, expected, rtol=0, atol=0)
|
|
assert ctx.sp_shard_metadata == ({} if shard_group is None else {shard_group: sequence_length})
|
|
assert ctx.sp_original_seq_len == (sequence_length if padding and shard_group is None else None)
|
|
|
|
|
|
@pytest.mark.parametrize("phase", ["success", "non_tensor", "forward_error", "split_error"])
|
|
@pytest.mark.parametrize("keyword", [False, True])
|
|
@torch.inference_mode()
|
|
def test_split_hook_does_not_retain_arguments_on_any_exit(parallel_state, phase, keyword):
|
|
class Module(nn.Module):
|
|
def forward(self, x):
|
|
if phase == "forward_error":
|
|
raise RuntimeError("forward failed")
|
|
return None if phase == "non_tensor" else x
|
|
|
|
module = Module()
|
|
index = 1 if phase == "split_error" else 0
|
|
hook = SequenceParallelSplitHook(
|
|
{index: SequenceParallelInput(split_dim=1, split_output=True, clone_shard=True)},
|
|
SequenceParallelConfig(ulysses_degree=4),
|
|
)
|
|
HookRegistry.get_or_create(module).register_hook("sp", hook)
|
|
refs = []
|
|
|
|
def run():
|
|
tensor = torch.randn(1, 16, 8)
|
|
refs.append(weakref.ref(tensor))
|
|
return module(x=tensor) if keyword else module(tensor)
|
|
|
|
with set_forward_context():
|
|
if phase.endswith("error"):
|
|
with pytest.raises((RuntimeError, ValueError)):
|
|
run()
|
|
else:
|
|
result = run()
|
|
assert (result is None) == (phase == "non_tensor")
|
|
assert all(ref() is None for ref in refs)
|
|
|
|
|
|
@pytest.mark.parametrize("keyword", [False, True])
|
|
@pytest.mark.parametrize("tensor_source", [False, True])
|
|
@torch.inference_mode()
|
|
def test_partial_output_split_keeps_only_lengths_across_calls(parallel_state, keyword, tensor_source):
|
|
parallel_state(rank=1)
|
|
|
|
class Module(nn.Module):
|
|
def forward(self, x, text):
|
|
return x
|
|
|
|
module = Module()
|
|
hook = SequenceParallelSplitHook(
|
|
{0: SequenceParallelPartialInput(split_dim=0, text_len_source="text", split_output=True)},
|
|
SequenceParallelConfig(ulysses_degree=4),
|
|
)
|
|
HookRegistry.get_or_create(module).register_hook("sp", hook)
|
|
refs = []
|
|
|
|
def run(text_length):
|
|
tensor = torch.arange(text_length + 8).reshape(-1, 1)
|
|
text = torch.zeros(text_length, 8) if tensor_source else text_length
|
|
refs.append(weakref.ref(tensor))
|
|
if tensor_source:
|
|
refs.append(weakref.ref(text))
|
|
return module(x=tensor, text=text) if keyword else module(tensor, text)
|
|
|
|
for text_length in (2, 4):
|
|
with set_forward_context():
|
|
result = run(text_length)
|
|
expected = torch.cat([torch.arange(text_length), torch.arange(text_length + 2, text_length + 4)]).reshape(-1, 1)
|
|
torch.testing.assert_close(result, expected)
|
|
assert all(ref() is None for ref in refs)
|
|
assert hook._text_len_cache == {"text": text_length}
|