118 lines
4.0 KiB
Python
118 lines
4.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""Equivalence test for DreamZero self-attn QKV fusion.
|
|
|
|
CausalWanSelfAttention now projects q/k/v with a single QKVParallelLinear and
|
|
splits the result, instead of three separate ColumnParallelLinear. The DreamZero
|
|
checkpoint still ships separate q/k/v weights, so load_weights routes each into
|
|
the packed `qkv` param with its shard id. This test proves the chain
|
|
"load separate q/k/v via shard id -> single GEMM -> split" reproduces the three
|
|
independent projections, given identical weights.
|
|
|
|
CPU-only: a single-rank (world_size=1) tensor-parallel group is initialized and
|
|
the GEMM dispatcher is forced to the default (CPU-compatible) path, following
|
|
tests/diffusion/layers/test_adalayernorm.py.
|
|
"""
|
|
|
|
import os
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _init_distributed():
|
|
"""Minimal world_size=1 TP group required to build QKVParallelLinear."""
|
|
from vllm.distributed.parallel_state import (
|
|
cleanup_dist_env_and_memory,
|
|
init_distributed_environment,
|
|
initialize_model_parallel,
|
|
)
|
|
|
|
os.environ.setdefault("MASTER_ADDR", "localhost")
|
|
os.environ.setdefault("MASTER_PORT", "29507")
|
|
init_distributed_environment(
|
|
world_size=1,
|
|
rank=0,
|
|
local_rank=0,
|
|
distributed_init_method="env://",
|
|
)
|
|
initialize_model_parallel()
|
|
yield
|
|
cleanup_dist_env_and_memory()
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _force_default_gemm(monkeypatch):
|
|
"""Force CPU-compatible GEMM dispatch for CPU test tensors."""
|
|
from vllm.model_executor.layers.utils import default_unquantized_gemm
|
|
|
|
monkeypatch.setattr(
|
|
"vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
|
|
lambda: default_unquantized_gemm,
|
|
)
|
|
|
|
|
|
def test_fused_qkv_load_and_split_matches_separate_projections():
|
|
from vllm.model_executor.layers.linear import QKVParallelLinear
|
|
|
|
torch.manual_seed(0)
|
|
dim, head_dim, num_heads = 16, 4, 4 # inner_dim == dim, no GQA
|
|
|
|
qkv = QKVParallelLinear(
|
|
hidden_size=dim,
|
|
head_size=head_dim,
|
|
total_num_heads=num_heads,
|
|
bias=True,
|
|
)
|
|
|
|
# Distinct random weight + bias per projection, loaded the exact way
|
|
# DreamZero.load_weights does: param.weight_loader(param, tensor, shard_id).
|
|
refs: dict[str, tuple[torch.Tensor, torch.Tensor]] = {}
|
|
for shard_id in ("q", "k", "v"):
|
|
w = torch.randn(dim, dim)
|
|
b = torch.randn(dim)
|
|
refs[shard_id] = (w, b)
|
|
qkv.weight.weight_loader(qkv.weight, w, shard_id)
|
|
qkv.bias.weight_loader(qkv.bias, b, shard_id)
|
|
|
|
x = torch.randn(2, 3, dim)
|
|
fused, _ = qkv(x)
|
|
q, k, v = fused.split([dim, dim, dim], dim=-1)
|
|
|
|
for got, shard_id in zip((q, k, v), ("q", "k", "v"), strict=True):
|
|
w, b = refs[shard_id]
|
|
ref = torch.nn.functional.linear(x, w, b)
|
|
torch.testing.assert_close(got, ref)
|
|
|
|
|
|
def test_fused_qkv_shard_ids_are_not_transposed():
|
|
# Guard against q/k/v being routed to the wrong slice: give each projection a
|
|
# constant, distinct weight so a swapped shard id would surface immediately.
|
|
from vllm.model_executor.layers.linear import QKVParallelLinear
|
|
|
|
dim, head_dim, num_heads = 16, 4, 4
|
|
qkv = QKVParallelLinear(
|
|
hidden_size=dim,
|
|
head_size=head_dim,
|
|
total_num_heads=num_heads,
|
|
bias=True,
|
|
)
|
|
|
|
scale = {"q": 1.0, "k": 2.0, "v": 3.0}
|
|
for shard_id, s in scale.items():
|
|
qkv.weight.weight_loader(qkv.weight, torch.full((dim, dim), s), shard_id)
|
|
qkv.bias.weight_loader(qkv.bias, torch.zeros(dim), shard_id)
|
|
|
|
x = torch.ones(1, dim)
|
|
fused, _ = qkv(x)
|
|
q, k, v = fused.split([dim, dim, dim], dim=-1)
|
|
|
|
# F.linear(ones, full(s)) == s * dim per element.
|
|
torch.testing.assert_close(q, torch.full((1, dim), scale["q"] * dim))
|
|
torch.testing.assert_close(k, torch.full((1, dim), scale["k"] * dim))
|
|
torch.testing.assert_close(v, torch.full((1, dim), scale["v"] * dim))
|