818 lines
30 KiB
Python
818 lines
30 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||
|
||
"""Tests for Ulysses + Ring sequence-parallel attention correctness.
|
||
|
||
What is tested
|
||
--------------
|
||
* ``test_sequence_parallel`` verifies that the ``Attention`` layer produces
|
||
numerically equivalent results whether the sequence is processed on a single
|
||
rank (baseline, SP=1) or sharded across multiple ranks via Ulysses/Ring SP.
|
||
The test spawns two separate multi-process runs with ``torch.multiprocessing.spawn``:
|
||
1. Baseline – world_size=1, ulysses_degree=1, ring_degree=1.
|
||
2. SP run – world_size=ulysses_degree*ring_degree, each rank holds a
|
||
contiguous slice of the full sequence.
|
||
After both runs, rank-0 output tensors are compared element-wise with a
|
||
tolerance appropriate for the dtype (bfloat16).
|
||
|
||
SP-plan hooks are NOT applied in this test
|
||
------------------------------------------
|
||
``ForwardContext.sp_plan_hooks_applied`` remains ``False``. As a result,
|
||
``ForwardContext.sp_active`` falls back to the "no hooks" branch: SP is
|
||
considered active whenever ``parallel_config.sequence_parallel_size > 1``.
|
||
This makes the test self-contained and suitable for standalone CI runs that
|
||
do not exercise the full model-registry pipeline.
|
||
"""
|
||
|
||
import os
|
||
import pickle
|
||
import tempfile
|
||
|
||
import pytest
|
||
import torch
|
||
|
||
from tests.helpers.mark import hardware_test
|
||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, QueryRange
|
||
from vllm_omni.diffusion.attention.layer import Attention
|
||
from vllm_omni.diffusion.attention.parallel.allgather_kv import (
|
||
AllGatherKVParallelAttention,
|
||
)
|
||
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.forward_context import set_forward_context
|
||
from vllm_omni.platforms import current_omni_platform
|
||
|
||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||
|
||
|
||
def update_environment_variables(envs_dict: dict[str, str]):
|
||
"""Update multiple environment variables with logging."""
|
||
for k, v in envs_dict.items():
|
||
os.environ[k] = v
|
||
|
||
|
||
def seed_everything(seed: int):
|
||
torch.manual_seed(seed)
|
||
current_omni_platform.manual_seed(seed)
|
||
|
||
|
||
class TestAttentionModel(torch.nn.Module):
|
||
"""Test model using Attention layer."""
|
||
|
||
def __init__(
|
||
self,
|
||
num_heads: int,
|
||
head_size: int,
|
||
hidden_size: int,
|
||
causal: bool = False,
|
||
num_kv_heads: int | None = None,
|
||
scatter_idx: int = 2,
|
||
gather_idx: int = 1,
|
||
use_sync: bool = False,
|
||
):
|
||
super().__init__()
|
||
self.num_heads = num_heads
|
||
self.head_size = head_size
|
||
self.hidden_size = hidden_size
|
||
self.attention = Attention(
|
||
num_heads=num_heads,
|
||
head_size=head_size,
|
||
causal=causal,
|
||
softmax_scale=1.0 / (head_size**0.5),
|
||
num_kv_heads=num_kv_heads,
|
||
scatter_idx=scatter_idx,
|
||
gather_idx=gather_idx,
|
||
use_sync=use_sync,
|
||
)
|
||
# Linear projection layers for Q, K, V
|
||
self.q_proj = torch.nn.Linear(hidden_size, num_heads * head_size)
|
||
self.k_proj = torch.nn.Linear(hidden_size, (num_kv_heads or num_heads) * head_size)
|
||
self.v_proj = torch.nn.Linear(hidden_size, (num_kv_heads or num_heads) * head_size)
|
||
self.o_proj = torch.nn.Linear(num_heads * head_size, hidden_size)
|
||
|
||
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||
"""Forward pass through attention layer."""
|
||
batch_size, seq_len, _ = hidden_states.shape
|
||
|
||
# Project to Q, K, V
|
||
q = self.q_proj(hidden_states)
|
||
k = self.k_proj(hidden_states)
|
||
v = self.v_proj(hidden_states)
|
||
|
||
# Reshape to (batch_size, seq_len, num_heads, head_size)
|
||
q = q.view(batch_size, seq_len, self.num_heads, self.head_size)
|
||
k = k.view(batch_size, seq_len, k.shape[-1] // self.head_size, self.head_size)
|
||
v = v.view(batch_size, seq_len, v.shape[-1] // self.head_size, self.head_size)
|
||
|
||
# Apply attention
|
||
attn_output = self.attention(q, k, v)
|
||
|
||
# Reshape back and project
|
||
attn_output = attn_output.view(batch_size, seq_len, -1)
|
||
output = self.o_proj(attn_output)
|
||
|
||
return output
|
||
|
||
|
||
class TestMultiLayerAttentionModel(torch.nn.Module):
|
||
"""Test model with multiple attention layers."""
|
||
|
||
def __init__(
|
||
self,
|
||
num_layers: int,
|
||
num_heads: int,
|
||
head_size: int,
|
||
hidden_size: int,
|
||
causal: bool = True,
|
||
num_kv_heads: int | None = None,
|
||
scatter_idx: int = 2,
|
||
gather_idx: int = 1,
|
||
use_sync: bool = False,
|
||
):
|
||
super().__init__()
|
||
self.num_layers = num_layers
|
||
self.layers = torch.nn.ModuleList(
|
||
[
|
||
TestAttentionModel(
|
||
num_heads=num_heads,
|
||
head_size=head_size,
|
||
hidden_size=hidden_size,
|
||
causal=causal,
|
||
num_kv_heads=num_kv_heads,
|
||
scatter_idx=scatter_idx,
|
||
gather_idx=gather_idx,
|
||
use_sync=use_sync,
|
||
)
|
||
for _ in range(num_layers)
|
||
]
|
||
)
|
||
|
||
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||
"""Forward pass through multiple attention layers."""
|
||
for layer in self.layers:
|
||
hidden_states = hidden_states + layer(hidden_states)
|
||
return hidden_states
|
||
|
||
|
||
class _MockAllGatherSPGroup:
|
||
def __init__(self, *, rank: int, gather_chunks: list[list[torch.Tensor]]) -> None:
|
||
self.allgather_group = object()
|
||
self.allgather_world_size = len(gather_chunks[0])
|
||
self.allgather_rank = rank
|
||
self._gather_chunks = list(gather_chunks)
|
||
self.gathered_input_shapes: list[tuple[int, ...]] = []
|
||
|
||
def all_gather(self, input_: torch.Tensor, dim: int = 0, separate_tensors: bool = False, group=None):
|
||
assert not separate_tensors
|
||
assert group is self.allgather_group
|
||
self.gathered_input_shapes.append(tuple(input_.shape))
|
||
chunks = self._gather_chunks.pop(0)
|
||
return torch.cat(chunks, dim=dim)
|
||
|
||
|
||
def test_allgather_kv_slices_full_dense_mask_to_local_query_rows():
|
||
rank = 1
|
||
joint_len = 1
|
||
img_seq_local = 2
|
||
img_seq_full = 4
|
||
key_chunks = [
|
||
torch.full((1, img_seq_local, 1, 1), 10.0),
|
||
torch.full((1, img_seq_local, 1, 1), 20.0),
|
||
]
|
||
value_chunks = [
|
||
torch.full((1, img_seq_local, 1, 1), 30.0),
|
||
torch.full((1, img_seq_local, 1, 1), 40.0),
|
||
]
|
||
strategy = AllGatherKVParallelAttention(
|
||
_MockAllGatherSPGroup(rank=rank, gather_chunks=[key_chunks, value_chunks]),
|
||
)
|
||
|
||
query = torch.zeros((1, img_seq_local, 1, 1))
|
||
key = key_chunks[rank]
|
||
value = value_chunks[rank]
|
||
joint = torch.ones((1, joint_len, 1, 1))
|
||
mask = torch.arange((joint_len + img_seq_full) * (joint_len + img_seq_full))
|
||
mask = (mask % 2 == 0).view(1, 1, joint_len + img_seq_full, joint_len + img_seq_full)
|
||
metadata = AttentionMetadata(
|
||
attn_mask=mask,
|
||
joint_query=joint,
|
||
joint_key=joint,
|
||
joint_value=joint,
|
||
joint_strategy="front",
|
||
)
|
||
|
||
q_local, k_full, v_full, metadata_local, _ = strategy.pre_attention(query, key, value, metadata)
|
||
|
||
expected_rows = torch.cat(
|
||
[
|
||
mask[..., :joint_len, :],
|
||
mask[
|
||
...,
|
||
joint_len + rank * img_seq_local : joint_len + (rank + 1) * img_seq_local,
|
||
:,
|
||
],
|
||
],
|
||
dim=-2,
|
||
)
|
||
assert q_local.shape[1] == joint_len + img_seq_local
|
||
assert k_full.shape[1] == joint_len + img_seq_full
|
||
assert v_full.shape[1] == joint_len + img_seq_full
|
||
torch.testing.assert_close(k_full[:, joint_len:], torch.cat(key_chunks, dim=1))
|
||
torch.testing.assert_close(v_full[:, joint_len:], torch.cat(value_chunks, dim=1))
|
||
assert metadata_local is not None
|
||
assert torch.equal(metadata_local.attn_mask, expected_rows)
|
||
|
||
|
||
def test_allgather_kv_slices_rear_joint_dense_mask():
|
||
rank = 1
|
||
joint_len = 1
|
||
img_seq_local = 2
|
||
img_seq_full = 4
|
||
key_chunks = [torch.zeros((1, img_seq_local, 1, 1)) for _ in range(2)]
|
||
value_chunks = [torch.zeros_like(chunk) for chunk in key_chunks]
|
||
strategy = AllGatherKVParallelAttention(
|
||
_MockAllGatherSPGroup(rank=rank, gather_chunks=[key_chunks, value_chunks]),
|
||
)
|
||
query = torch.zeros((1, img_seq_local, 1, 1))
|
||
joint = torch.ones((1, joint_len, 1, 1))
|
||
mask = torch.arange((img_seq_full + joint_len) ** 2)
|
||
mask = (mask % 2 == 0).view(1, 1, img_seq_full + joint_len, img_seq_full + joint_len)
|
||
metadata = AttentionMetadata(
|
||
attn_mask=mask,
|
||
joint_query=joint,
|
||
joint_key=joint,
|
||
joint_value=joint,
|
||
joint_strategy="rear",
|
||
)
|
||
|
||
q_local, _, _, metadata_local, _ = strategy.pre_attention(query, key_chunks[rank], value_chunks[rank], metadata)
|
||
|
||
img_start = rank * img_seq_local
|
||
expected_rows = torch.cat(
|
||
[
|
||
mask[..., img_start : img_start + img_seq_local, :],
|
||
mask[..., img_seq_full : img_seq_full + joint_len, :],
|
||
],
|
||
dim=-2,
|
||
)
|
||
assert q_local.shape[1] == img_seq_local + joint_len
|
||
assert metadata_local is not None
|
||
assert torch.equal(metadata_local.attn_mask, expected_rows)
|
||
|
||
|
||
def test_allgather_kv_rejects_invalid_joint_strategy():
|
||
chunks = [[torch.zeros((1, 2, 1, 1)) for _ in range(2)] for _ in range(2)]
|
||
strategy = AllGatherKVParallelAttention(_MockAllGatherSPGroup(rank=0, gather_chunks=chunks))
|
||
joint = torch.ones((1, 1, 1, 1))
|
||
metadata = AttentionMetadata(
|
||
joint_query=joint,
|
||
joint_key=joint,
|
||
joint_value=joint,
|
||
joint_strategy="back",
|
||
)
|
||
|
||
with pytest.raises(ValueError, match="Unsupported joint_strategy"):
|
||
strategy.pre_attention(torch.zeros((1, 2, 1, 1)), chunks[0][0], chunks[1][0], metadata)
|
||
|
||
|
||
def test_allgather_kv_preserves_global_spans_and_sets_query_ranges():
|
||
strategy = AllGatherKVParallelAttention(
|
||
_MockAllGatherSPGroup(
|
||
rank=1,
|
||
gather_chunks=[
|
||
[torch.zeros((1, 2, 1, 1)), torch.zeros((1, 2, 1, 1))],
|
||
[torch.zeros((1, 2, 1, 1)), torch.zeros((1, 2, 1, 1))],
|
||
],
|
||
),
|
||
)
|
||
query = torch.zeros((1, 2, 1, 1))
|
||
key = torch.zeros((1, 2, 1, 1))
|
||
value = torch.zeros((1, 2, 1, 1))
|
||
joint = torch.ones((1, 1, 1, 1))
|
||
metadata = AttentionMetadata(
|
||
joint_query=joint,
|
||
joint_key=joint,
|
||
joint_value=joint,
|
||
full_attn_spans=[
|
||
[
|
||
(0, 1), # joint span: kept on every rank.
|
||
(2, 5), # image span crossing rank 0 and rank 1 image shards.
|
||
]
|
||
],
|
||
)
|
||
|
||
_, _, _, metadata_out, _ = strategy.pre_attention(query, key, value, metadata)
|
||
|
||
assert metadata_out is not None
|
||
assert metadata_out.full_attn_spans == [[(0, 1), (2, 5)]]
|
||
assert metadata_out.query_ranges == (
|
||
QueryRange(0, 1, 0),
|
||
QueryRange(1, 3, 3),
|
||
)
|
||
|
||
|
||
def test_allgather_kv_query_ranges_include_reused_prefix_offset():
|
||
strategy = AllGatherKVParallelAttention(
|
||
_MockAllGatherSPGroup(
|
||
rank=1,
|
||
gather_chunks=[
|
||
[torch.zeros((1, 2, 1, 1)), torch.zeros((1, 2, 1, 1))],
|
||
[torch.zeros((1, 2, 1, 1)), torch.zeros((1, 2, 1, 1))],
|
||
],
|
||
),
|
||
)
|
||
query = torch.zeros((1, 2, 1, 1))
|
||
key = torch.zeros((1, 2, 1, 1))
|
||
value = torch.zeros((1, 2, 1, 1))
|
||
joint_query = torch.ones((1, 1, 1, 1))
|
||
# K/V retain two reused AR-prefix tokens that have already been removed
|
||
# from Q and from the attention mask's query rows.
|
||
joint_key = torch.ones((1, 3, 1, 1))
|
||
joint_value = torch.ones((1, 3, 1, 1))
|
||
mask = torch.arange(5 * 7).view(1, 1, 5, 7)
|
||
metadata = AttentionMetadata(
|
||
attn_mask=mask,
|
||
joint_query=joint_query,
|
||
joint_key=joint_key,
|
||
joint_value=joint_value,
|
||
full_attn_spans=[[(0, 7)]],
|
||
)
|
||
|
||
_, _, _, metadata_out, _ = strategy.pre_attention(query, key, value, metadata)
|
||
|
||
assert metadata_out is not None
|
||
assert metadata_out.query_ranges == (
|
||
QueryRange(0, 1, 2),
|
||
QueryRange(1, 3, 5),
|
||
)
|
||
expected_mask = torch.cat([mask[..., :1, :], mask[..., 3:5, :]], dim=-2)
|
||
assert torch.equal(metadata_out.attn_mask, expected_mask)
|
||
|
||
|
||
def test_allgather_kv_allows_empty_full_attn_spans():
|
||
strategy = AllGatherKVParallelAttention(
|
||
_MockAllGatherSPGroup(
|
||
rank=0,
|
||
gather_chunks=[
|
||
[torch.zeros((1, 2, 1, 1)), torch.zeros((1, 2, 1, 1))],
|
||
[torch.zeros((1, 2, 1, 1)), torch.zeros((1, 2, 1, 1))],
|
||
],
|
||
),
|
||
)
|
||
query = torch.zeros((1, 2, 1, 1))
|
||
key = torch.zeros((1, 2, 1, 1))
|
||
value = torch.zeros((1, 2, 1, 1))
|
||
metadata = AttentionMetadata(full_attn_spans=[[]])
|
||
|
||
_, _, _, metadata_out, _ = strategy.pre_attention(query, key, value, metadata)
|
||
|
||
assert metadata_out is not None
|
||
assert metadata_out.full_attn_spans == [[]]
|
||
assert metadata_out.query_ranges == (QueryRange(0, 2, 0),)
|
||
|
||
|
||
def test_allgather_kv_keeps_gathered_kv_compressed_for_gqa():
|
||
rank = 0
|
||
img_seq_local = 2
|
||
kv_heads = 2
|
||
repeat_num = 2
|
||
q_heads = kv_heads * repeat_num
|
||
key_chunks = [
|
||
torch.arange(1 * img_seq_local * kv_heads * 1, dtype=torch.float32).view(1, img_seq_local, kv_heads, 1),
|
||
torch.arange(100, 100 + 1 * img_seq_local * kv_heads * 1, dtype=torch.float32).view(
|
||
1, img_seq_local, kv_heads, 1
|
||
),
|
||
]
|
||
value_chunks = [chunk + 1000 for chunk in key_chunks]
|
||
sp_group = _MockAllGatherSPGroup(rank=rank, gather_chunks=[key_chunks, value_chunks])
|
||
strategy = AllGatherKVParallelAttention(sp_group)
|
||
|
||
query = torch.zeros((1, img_seq_local, q_heads, 1))
|
||
_, k_full, v_full, _, _ = strategy.pre_attention(query, key_chunks[rank], value_chunks[rank], AttentionMetadata())
|
||
|
||
assert sp_group.gathered_input_shapes == [
|
||
(1, img_seq_local, kv_heads, 1),
|
||
(1, img_seq_local, kv_heads, 1),
|
||
]
|
||
assert k_full.shape == (1, img_seq_local * 2, kv_heads, 1)
|
||
assert v_full.shape == (1, img_seq_local * 2, kv_heads, 1)
|
||
expected_key = torch.cat(key_chunks, dim=1)
|
||
expected_value = torch.cat(value_chunks, dim=1)
|
||
torch.testing.assert_close(k_full, expected_key)
|
||
torch.testing.assert_close(v_full, expected_value)
|
||
|
||
|
||
@hardware_test(res={"cuda": "L4"}, num_cards=4)
|
||
@pytest.mark.parametrize(
|
||
"test_model_cls",
|
||
[
|
||
TestMultiLayerAttentionModel,
|
||
],
|
||
)
|
||
@pytest.mark.parametrize(
|
||
("ulysses_degree", "ring_degree", "allgather_degree", "num_kv_heads"),
|
||
[
|
||
pytest.param(2, 2, 1, None, id="ulysses-ring"),
|
||
pytest.param(1, 1, 2, None, id="allgather-kv"),
|
||
pytest.param(1, 2, 1, 2, id="ring-gqa"),
|
||
pytest.param(1, 2, 1, 1, id="ring-mqa"),
|
||
],
|
||
)
|
||
@pytest.mark.parametrize("batch_size", [2])
|
||
@pytest.mark.parametrize("seq_len", [16])
|
||
@pytest.mark.parametrize("num_heads", [8])
|
||
@pytest.mark.parametrize("head_size", [8])
|
||
@pytest.mark.parametrize("causal", [False])
|
||
@pytest.mark.parametrize("dtype", [torch.bfloat16])
|
||
@pytest.mark.parametrize("use_sync", [False])
|
||
@pytest.mark.parametrize("dynamic", [False])
|
||
@pytest.mark.parametrize("use_compile", [False])
|
||
def test_sequence_parallel(
|
||
ulysses_degree: int,
|
||
ring_degree: int,
|
||
allgather_degree: int,
|
||
test_model_cls: type[torch.nn.Module],
|
||
dtype: torch.dtype,
|
||
causal: bool,
|
||
use_sync: bool,
|
||
dynamic: bool,
|
||
use_compile: bool,
|
||
batch_size: int,
|
||
seq_len: int,
|
||
num_heads: int,
|
||
head_size: int,
|
||
num_kv_heads: int | None,
|
||
):
|
||
"""Compare Ulysses/Ring and AllGather-KV SP against a single-rank run."""
|
||
sequence_parallel_size = allgather_degree if allgather_degree > 1 else ulysses_degree * ring_degree
|
||
|
||
# Skip if not enough GPUs available
|
||
available_gpus = current_omni_platform.get_device_count()
|
||
if available_gpus < sequence_parallel_size:
|
||
pytest.skip(f"Test requires {sequence_parallel_size} GPUs but only {available_gpus} available")
|
||
|
||
# Create temporary files to share results between processes
|
||
with tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") as f:
|
||
baseline_output_file = f.name
|
||
with tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") as f:
|
||
sp_output_file = f.name
|
||
with tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") as f:
|
||
model_state_file = f.name
|
||
with tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") as f:
|
||
input_data_file = f.name
|
||
|
||
try:
|
||
# Step 1: Run without SP (baseline with ulysses_degree=1, ring_degree=1)
|
||
print("\n[Baseline] Running without SP (ulysses_degree=1, ring_degree=1)...")
|
||
torch.multiprocessing.spawn(
|
||
ulysses_attention_on_test_model,
|
||
args=(
|
||
1, # num_processes = 1 for baseline
|
||
test_model_cls,
|
||
batch_size,
|
||
seq_len,
|
||
num_heads,
|
||
head_size,
|
||
num_kv_heads,
|
||
dtype,
|
||
causal,
|
||
use_sync,
|
||
dynamic,
|
||
use_compile,
|
||
1, # ulysses_degree = 1
|
||
1, # ring_degree = 1
|
||
1, # allgather_degree = 1 (baseline: no AllGather-KV SP)
|
||
1, # sequence_parallel_size = 1
|
||
baseline_output_file,
|
||
model_state_file,
|
||
input_data_file,
|
||
True, # is_baseline
|
||
),
|
||
nprocs=1,
|
||
)
|
||
|
||
# Step 2: Run with SP enabled
|
||
print(
|
||
"\n[SP Test] Running with SP "
|
||
f"(ulysses={ulysses_degree}, ring={ring_degree}, allgather={allgather_degree})..."
|
||
)
|
||
torch.multiprocessing.spawn(
|
||
ulysses_attention_on_test_model,
|
||
args=(
|
||
sequence_parallel_size, # num_processes
|
||
test_model_cls,
|
||
batch_size,
|
||
seq_len,
|
||
num_heads,
|
||
head_size,
|
||
num_kv_heads,
|
||
dtype,
|
||
causal,
|
||
use_sync,
|
||
dynamic,
|
||
use_compile,
|
||
ulysses_degree,
|
||
ring_degree,
|
||
allgather_degree,
|
||
sequence_parallel_size,
|
||
sp_output_file,
|
||
model_state_file,
|
||
input_data_file,
|
||
False, # is_baseline
|
||
),
|
||
nprocs=sequence_parallel_size,
|
||
)
|
||
|
||
# Step 3: Verify input consistency and compare outputs
|
||
print(f"\n{'=' * 80}")
|
||
print("Verifying input data consistency...")
|
||
with open(input_data_file, "rb") as f:
|
||
input_data = pickle.load(f)
|
||
input_checksum = hash(input_data.tobytes())
|
||
print(f" Input data shape: {input_data.shape}")
|
||
print(f" Input data checksum: {input_checksum}")
|
||
print(" ✓ Both baseline and SP used the same input data")
|
||
|
||
print(f"\n{'=' * 80}")
|
||
print("Comparing outputs between baseline and SP...")
|
||
with open(baseline_output_file, "rb") as f:
|
||
baseline_output = pickle.load(f)
|
||
with open(sp_output_file, "rb") as f:
|
||
sp_output = pickle.load(f)
|
||
|
||
# Convert to tensors for comparison
|
||
baseline_tensor = torch.tensor(baseline_output)
|
||
sp_tensor = torch.tensor(sp_output)
|
||
|
||
print(f" Baseline output shape: {baseline_tensor.shape}")
|
||
print(f" SP output shape: {sp_tensor.shape}")
|
||
assert baseline_tensor.shape == sp_tensor.shape, "Output shapes must match!"
|
||
|
||
# Calculate differences
|
||
abs_diff = torch.abs(baseline_tensor - sp_tensor)
|
||
max_abs_diff = abs_diff.max().item()
|
||
mean_abs_diff = abs_diff.mean().item()
|
||
|
||
# Calculate relative difference (avoid division by zero)
|
||
baseline_abs = torch.abs(baseline_tensor)
|
||
relative_diff = abs_diff / (baseline_abs + 1e-8)
|
||
max_relative_diff = relative_diff.max().item()
|
||
mean_relative_diff = relative_diff.mean().item()
|
||
|
||
print(f"\n{'=' * 80}")
|
||
print("Output Difference Analysis:")
|
||
print(f" - Max absolute difference: {max_abs_diff:.6e}")
|
||
print(f" - Mean absolute difference: {mean_abs_diff:.6e}")
|
||
print(f" - Max relative difference: {max_relative_diff:.6e}")
|
||
print(f" - Mean relative difference: {mean_relative_diff:.6e}")
|
||
print(f" - Baseline output range: [{baseline_tensor.min().item():.6e}, {baseline_tensor.max().item():.6e}]")
|
||
print(f" - SP output range: [{sp_tensor.min().item():.6e}, {sp_tensor.max().item():.6e}]")
|
||
print(f"{'=' * 80}\n")
|
||
|
||
# Assert that differences are within acceptable tolerance
|
||
# For FP16/BF16, we expect some numerical differences due to different computation order under parallelism.
|
||
# If we use the same backend (e.g. Flash Attention) for both baseline and SP, differences should be smaller.
|
||
if dtype == torch.float16:
|
||
atol, rtol = 5e-2, 5e-2 # Increased tolerance for Ring Attention
|
||
elif dtype == torch.bfloat16:
|
||
atol, rtol = 5e-2, 5e-2 # Increased tolerance for Ring Attention
|
||
else:
|
||
atol, rtol = 1e-5, 1e-4
|
||
|
||
assert max_abs_diff < atol or max_relative_diff < rtol, (
|
||
f"Output difference too large: max_abs_diff={max_abs_diff:.6e}, "
|
||
f"max_relative_diff={max_relative_diff:.6e}, "
|
||
f"tolerance: atol={atol}, rtol={rtol}"
|
||
)
|
||
|
||
print("✓ Test passed: SP output matches baseline within tolerance")
|
||
|
||
finally:
|
||
# Clean up temporary files
|
||
for f in [baseline_output_file, sp_output_file, model_state_file, input_data_file]:
|
||
if os.path.exists(f):
|
||
os.remove(f)
|
||
|
||
|
||
def ulysses_attention_on_test_model(
|
||
local_rank: int,
|
||
world_size: int,
|
||
test_model_cls: type[torch.nn.Module],
|
||
batch_size: int,
|
||
seq_len: int,
|
||
num_heads: int,
|
||
head_size: int,
|
||
num_kv_heads: int | None,
|
||
dtype: torch.dtype,
|
||
causal: bool,
|
||
use_sync: bool,
|
||
dynamic: bool,
|
||
use_compile: bool,
|
||
ulysses_degree: int,
|
||
ring_degree: int,
|
||
allgather_degree: int,
|
||
sequence_parallel_size: int,
|
||
output_file: str,
|
||
model_state_file: str,
|
||
input_data_file: str,
|
||
is_baseline: bool,
|
||
):
|
||
"""Run Ulysses attention test on a test model and save results for comparison."""
|
||
# Use fixed seed for reproducibility across baseline and SP runs
|
||
RANDOM_SEED = 42
|
||
seed_everything(RANDOM_SEED)
|
||
|
||
if allgather_degree > 1:
|
||
sp_kind = f"allgather={allgather_degree}"
|
||
else:
|
||
sp_kind = f"ulysses={ulysses_degree}, ring={ring_degree}"
|
||
mode_str = "Baseline (no SP)" if is_baseline else f"SP ({sp_kind})"
|
||
print(f"\n[{mode_str}] Rank {local_rank}/{world_size} - Random seed set to {RANDOM_SEED}")
|
||
|
||
device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
|
||
current_omni_platform.set_device(device)
|
||
torch.set_default_device(device)
|
||
torch.set_default_dtype(dtype)
|
||
|
||
update_environment_variables(
|
||
{
|
||
"RANK": str(local_rank),
|
||
"LOCAL_RANK": str(local_rank),
|
||
"WORLD_SIZE": str(world_size),
|
||
"MASTER_ADDR": "localhost",
|
||
"MASTER_PORT": "12345",
|
||
}
|
||
)
|
||
# Initialize distributed environment
|
||
init_distributed_environment()
|
||
|
||
# Set up OmniDiffusionConfig with parallel config
|
||
parallel_config = DiffusionParallelConfig(
|
||
pipeline_parallel_size=1,
|
||
data_parallel_size=1,
|
||
tensor_parallel_size=1,
|
||
sequence_parallel_size=sequence_parallel_size,
|
||
ulysses_degree=ulysses_degree,
|
||
ring_degree=ring_degree,
|
||
allgather_degree=allgather_degree,
|
||
cfg_parallel_size=1,
|
||
)
|
||
|
||
od_config = OmniDiffusionConfig.from_kwargs(
|
||
model="test_model",
|
||
dtype=dtype,
|
||
parallel_config=parallel_config,
|
||
# This regression targets pytorch_attn_forward(). Do not let an
|
||
# installed FA/FA3 backend silently bypass the SDPA Ring path.
|
||
diffusion_attention_backend="TORCH_SDPA",
|
||
)
|
||
|
||
# Initialize model parallel
|
||
initialize_model_parallel(
|
||
data_parallel_size=1,
|
||
cfg_parallel_size=1,
|
||
sequence_parallel_size=sequence_parallel_size,
|
||
ulysses_degree=ulysses_degree,
|
||
ring_degree=ring_degree,
|
||
allgather_degree=allgather_degree,
|
||
tensor_parallel_size=1,
|
||
pipeline_parallel_size=1,
|
||
)
|
||
|
||
# Set the config so Attention can access it during init and forward
|
||
with set_forward_context(omni_diffusion_config=od_config), set_current_diffusion_config(od_config):
|
||
# Create model
|
||
hidden_size = num_heads * head_size
|
||
|
||
# Create model with appropriate parameters
|
||
model_kwargs = {
|
||
"num_heads": num_heads,
|
||
"head_size": head_size,
|
||
"hidden_size": hidden_size,
|
||
"causal": causal,
|
||
"num_kv_heads": num_kv_heads,
|
||
"scatter_idx": 2,
|
||
"gather_idx": 1,
|
||
"use_sync": use_sync,
|
||
}
|
||
|
||
if test_model_cls == TestMultiLayerAttentionModel:
|
||
model_kwargs["num_layers"] = 2
|
||
|
||
model = test_model_cls(**model_kwargs)
|
||
model = model.to(device).to(dtype)
|
||
|
||
# For baseline: Generate and save model state and input data
|
||
# This ensures both baseline and SP use exactly the same initialization
|
||
if is_baseline and local_rank == 0:
|
||
# Save model state for reuse (before any computation)
|
||
model_state = {k: v.cpu() for k, v in model.state_dict().items()}
|
||
with open(model_state_file, "wb") as f:
|
||
pickle.dump(model_state, f)
|
||
|
||
full_hidden_states = torch.randn(
|
||
(batch_size, seq_len, hidden_size),
|
||
dtype=dtype,
|
||
device="cpu",
|
||
)
|
||
with open(input_data_file, "wb") as f:
|
||
pickle.dump(full_hidden_states.detach().cpu().float().numpy(), f)
|
||
|
||
print("[Baseline] Saved model state and input data")
|
||
|
||
# Synchronize to ensure baseline has saved data before SP loads it
|
||
if world_size > 1:
|
||
torch.distributed.barrier()
|
||
|
||
# IMPORTANT: Both baseline and SP load the same model state and input data
|
||
# This ensures exact same initialization and input for fair comparison
|
||
with open(model_state_file, "rb") as f:
|
||
model_state = pickle.load(f)
|
||
model.load_state_dict({k: v.to(device).to(dtype) for k, v in model_state.items()})
|
||
|
||
with open(input_data_file, "rb") as f:
|
||
full_hidden_states_np = pickle.load(f)
|
||
full_hidden_states = torch.from_numpy(full_hidden_states_np).to(device).to(dtype)
|
||
|
||
print(f"[Rank {local_rank}] Loaded model state and full input data with shape {full_hidden_states.shape}")
|
||
|
||
# Split input sequence according to sequence parallel BEFORE model forward
|
||
# Each rank gets a contiguous chunk of the sequence dimension
|
||
local_seq_len = seq_len // sequence_parallel_size
|
||
start_idx = local_rank * local_seq_len
|
||
end_idx = start_idx + local_seq_len
|
||
hidden_states = full_hidden_states[:, start_idx:end_idx, :].contiguous()
|
||
|
||
print(
|
||
f"[Rank {local_rank}] Split input: local_seq_len={local_seq_len}, "
|
||
f"indices=[{start_idx}:{end_idx}], local_shape={hidden_states.shape}"
|
||
)
|
||
|
||
if dynamic:
|
||
torch._dynamo.mark_dynamic(hidden_states, 0)
|
||
torch._dynamo.mark_dynamic(hidden_states, 1)
|
||
|
||
# Compile model if requested
|
||
if use_compile:
|
||
model = torch.compile(model)
|
||
|
||
# Run forward pass with local sequence chunk
|
||
print(f"[Rank {local_rank}] Running forward pass...")
|
||
output = model(hidden_states)
|
||
print(f"[Rank {local_rank}] Forward pass completed, output shape: {output.shape}")
|
||
|
||
# Verify output shape
|
||
assert output.shape == (batch_size, local_seq_len, hidden_size), (
|
||
f"Output shape mismatch: expected {(batch_size, local_seq_len, hidden_size)}, got {output.shape}"
|
||
)
|
||
|
||
# Gather outputs from all ranks AFTER computation
|
||
if world_size > 1:
|
||
print(f"[Rank {local_rank}] Gathering outputs from all {world_size} ranks...")
|
||
# Gather all outputs to rank 0
|
||
gathered_outputs = [torch.zeros_like(output) for _ in range(world_size)]
|
||
torch.distributed.all_gather(gathered_outputs, output)
|
||
if local_rank == 0:
|
||
# Concatenate along sequence dimension to reconstruct full sequence
|
||
full_output = torch.cat(gathered_outputs, dim=1)
|
||
print(f"[Rank 0] Gathered and concatenated outputs: {full_output.shape}")
|
||
# Verify the full output shape matches expected
|
||
assert full_output.shape == (batch_size, seq_len, hidden_size), (
|
||
f"Gathered output shape mismatch: expected {(batch_size, seq_len, hidden_size)}, "
|
||
f"got {full_output.shape}"
|
||
)
|
||
else:
|
||
full_output = None
|
||
else:
|
||
# For baseline (world_size=1), output is already complete
|
||
full_output = output
|
||
print(f"[Rank 0] No gather needed (world_size=1), output shape: {full_output.shape}")
|
||
|
||
# Save output from rank 0 for comparison
|
||
if local_rank == 0:
|
||
output_np = full_output.detach().cpu().float().numpy()
|
||
with open(output_file, "wb") as f:
|
||
pickle.dump(output_np, f)
|
||
|
||
if allgather_degree > 1:
|
||
sp_kind = f"allgather={allgather_degree}"
|
||
else:
|
||
sp_kind = f"ulysses={ulysses_degree}, ring={ring_degree}"
|
||
mode_str = "baseline (no SP)" if is_baseline else f"SP ({sp_kind})"
|
||
print(
|
||
f"\n[{mode_str}] ✓ Saved output with shape {full_output.shape}:\n"
|
||
f" - batch_size={batch_size}, seq_len={seq_len}\n"
|
||
f" - num_heads={num_heads}, head_size={head_size}\n"
|
||
f" - dtype={dtype}, causal={causal}, use_sync={use_sync}\n"
|
||
)
|
||
|
||
destroy_distributed_env()
|