tzhouamandClaude Opus 5
9a399adba8
Scoped pre-commit over the 27 files this rebase touches: ruff format reflow, the vLLM-Omni SPDX copyright line on the new modules, and wording the typos hook rejected. The Python 3.10 mypy hook still reports 24 errors across five of these files. All 24 sit on pre-existing lines; none fall on a line this rebase changed, which the repository documents as the CI-skipped condition. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0146Xu5Sgj4vXxLxvLjV5YPj Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
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-Omni 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 *_args, **_kwargs: 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))
|