425 lines
15 KiB
Python
425 lines
15 KiB
Python
#
|
|
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
|
|
# Copyright 2023 The vLLM team.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
"""Unit tests for VWN-Eagle3 model components.
|
|
|
|
Tests cover PreVwnLayerV1, VwnLlamaDecoderLayer, VwnLlamaModel, and
|
|
Eagle3VwnLlamaForCausalLM using CPU-only execution with mocked VllmConfig.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from contextlib import contextmanager
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from vllm.config import CacheConfig, CompilationMode, VllmConfig, set_current_vllm_config
|
|
|
|
from vllm_ascend.ascend_config import init_ascend_config
|
|
from vllm_ascend.models.llama_eagle3_vwn import (
|
|
Eagle3VwnLlamaForCausalLM,
|
|
PreVwnLayerV1,
|
|
VwnLlamaDecoderLayer,
|
|
VwnLlamaModel,
|
|
)
|
|
|
|
_HIDDEN = 2048
|
|
_INTERMEDIATE = 6144
|
|
_VOCAB = 151936
|
|
_DRAFT_VOCAB = 35000
|
|
_NUM_HEADS = 32
|
|
_NUM_KV_HEADS = 4
|
|
_RMS_EPS = 1e-6
|
|
|
|
|
|
class _PassthroughAttn(nn.Module):
|
|
"""Replaces self_attn for CPU tests — returns input unchanged."""
|
|
|
|
def forward(self, *, positions, hidden_states, **kwargs):
|
|
return hidden_states
|
|
|
|
|
|
class _PassthroughMLP(nn.Module):
|
|
"""Replaces mlp for CPU tests — returns input unchanged."""
|
|
|
|
def forward(self, hidden_states):
|
|
return hidden_states
|
|
|
|
|
|
class _MockTPGroup:
|
|
"""Minimal mock for get_tp_group() when TP=1."""
|
|
|
|
rank_in_group = 0
|
|
world_size = 1
|
|
|
|
def all_reduce(self, *args, **kwargs):
|
|
pass
|
|
|
|
def all_gather(self, x, *args, **kwargs):
|
|
return x.unsqueeze(0)
|
|
|
|
def reduce_scatter(self, x, *args, **kwargs):
|
|
return x
|
|
|
|
|
|
def _mock_npu_ops_on_layer(layer):
|
|
"""Replace self_attn and mlp with passthrough modules for CPU testing."""
|
|
layer.self_attn = _PassthroughAttn()
|
|
layer.mlp = _PassthroughMLP()
|
|
|
|
|
|
def _cpu_rms_norm(x, weight, eps):
|
|
"""CPU fallback for torch_npu.npu_rms_norm.
|
|
|
|
Returns the normalized tensor and a placeholder rstd (None), matching the
|
|
2-tuple shape the production op yields so callers can unpack it.
|
|
"""
|
|
orig_dtype = x.dtype
|
|
x32 = x.float()
|
|
var = x32.pow(2).mean(-1, keepdim=True)
|
|
x32 = x32 * torch.rsqrt(var + eps)
|
|
out = (x32 * weight.float()).to(orig_dtype)
|
|
return out, None
|
|
|
|
|
|
def _cpu_add_rms_norm(x, residual, weight, eps):
|
|
"""CPU fallback for torch_npu.npu_add_rms_norm (returns 3-tuple)."""
|
|
x_plus_res = x + residual
|
|
out, _ = _cpu_rms_norm(x_plus_res, weight, eps)
|
|
return out, None, x_plus_res
|
|
|
|
|
|
def _cpu_add_rms_norm_bias(x, residual, weight, bias, eps):
|
|
"""CPU fallback for torch.ops._C_ascend.npu_add_rms_norm_bias."""
|
|
out, _, new_residual = _cpu_add_rms_norm(x, residual, weight, eps)
|
|
if bias is not None:
|
|
out = out + bias
|
|
return out, _, new_residual
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _mock_npu_env():
|
|
"""Patch TP group, Ascend config, and NPU ops so all tests run on CPU.
|
|
|
|
conftest.py stubs ``torch_npu.npu_rms_norm`` with a bare ``MagicMock()``,
|
|
which yields an empty iterator and breaks tuple unpacking. We override it
|
|
here (plus the related add/bias variants and the weight-prefetch hook) with
|
|
pure-torch CPU implementations so AscendRMSNorm can run on CPU runners.
|
|
"""
|
|
import torch_npu
|
|
|
|
_mock = _MockTPGroup()
|
|
mock_cfg = MagicMock()
|
|
mock_cfg.enable_flashcomm2_parallel_size = 0
|
|
mock_cfg.enable_context_parallel = False
|
|
mock_cfg.enable_flashcomm1 = False
|
|
mock_cfg.enable_matmul_allreduce = False
|
|
mock_cfg.weight_nz_mode = 1
|
|
mock_cfg.enable_mlapo = True
|
|
mock_cfg.enable_fused_mc2 = 0
|
|
mock_cfg.msmonitor_use_daemon = False
|
|
mock_cfg.enable_transpose_kv_cache_by_block = True
|
|
mock_cfg.finegrained_tp_config = MagicMock(
|
|
lmhead_tensor_parallel_size=0,
|
|
embedding_tensor_parallel_size=0,
|
|
oproj_tensor_parallel_size=0,
|
|
olora_tensor_parallel_size=0,
|
|
mlp_tensor_parallel_size=0,
|
|
)
|
|
|
|
_prefetch_mock = MagicMock()
|
|
|
|
with (
|
|
patch("vllm_ascend.ops.linear_op.get_tp_group", return_value=_mock),
|
|
patch("vllm.distributed.parallel_state.get_tp_group", return_value=_mock),
|
|
patch("vllm_ascend.ops.vocab_parallel_embedding.get_tp_group", return_value=_mock),
|
|
patch("vllm_ascend.utils.get_ascend_config", return_value=mock_cfg),
|
|
patch.object(torch.ops.vllm, "unquantized_gemm", F.linear),
|
|
patch.object(torch.ops.vllm, "maybe_calc_kv_scales", lambda *a, **kw: None),
|
|
patch.object(torch.ops.vllm, "maybe_pad_and_reduce", lambda x, *a, **kw: x),
|
|
patch("vllm.model_executor.layers.logits_processor.tensor_model_parallel_all_gather", lambda x, *a, **kw: x),
|
|
patch.object(torch_npu, "npu_rms_norm", side_effect=_cpu_rms_norm, create=True),
|
|
patch.object(torch_npu, "npu_add_rms_norm", side_effect=_cpu_add_rms_norm, create=True),
|
|
patch.object(
|
|
torch.ops._C_ascend,
|
|
"npu_add_rms_norm_bias",
|
|
side_effect=_cpu_add_rms_norm_bias,
|
|
create=True,
|
|
),
|
|
patch("vllm_ascend.ops.layernorm.get_weight_prefetch_method", return_value=_prefetch_mock),
|
|
# enable_cp() reads parallel_config.*_context_parallel_size and runs `> 1`.
|
|
# On MagicMock these fields yield TypeError on Python 3.12, so short-circuit
|
|
# the check everywhere it's imported.
|
|
patch("vllm_ascend.attention.attention_v1.enable_cp", return_value=False),
|
|
patch("vllm_ascend.attention.sfa_v1.enable_cp", return_value=False, create=True),
|
|
patch("vllm_ascend.attention.mla_v1.enable_cp", return_value=False, create=True),
|
|
):
|
|
yield
|
|
|
|
|
|
def _make_hf_config(
|
|
hidden_size=_HIDDEN,
|
|
vwn_m=4,
|
|
vwn_r=1.5,
|
|
num_hidden_layers=1,
|
|
draft_vocab_size=_DRAFT_VOCAB,
|
|
**extra,
|
|
):
|
|
"""Create a real LlamaConfig with VWN attributes.
|
|
|
|
Using a real config object avoids whack-a-mole with missing attributes
|
|
that LlamaDecoderLayer's deep init chain expects.
|
|
"""
|
|
from transformers import LlamaConfig
|
|
|
|
cfg = LlamaConfig(
|
|
hidden_size=hidden_size,
|
|
intermediate_size=_INTERMEDIATE,
|
|
num_attention_heads=_NUM_HEADS,
|
|
num_key_value_heads=_NUM_KV_HEADS,
|
|
num_hidden_layers=num_hidden_layers,
|
|
vocab_size=_VOCAB,
|
|
rms_norm_eps=_RMS_EPS,
|
|
max_position_embeddings=40960,
|
|
)
|
|
cfg.vwn_m = vwn_m
|
|
cfg.vwn_r = vwn_r
|
|
cfg.draft_vocab_size = draft_vocab_size
|
|
for k, v in extra.items():
|
|
setattr(cfg, k, v)
|
|
return cfg
|
|
|
|
|
|
def _create_vllm_config_for_vwn(
|
|
vwn_m=4,
|
|
vwn_r=1.5,
|
|
hidden_size=_HIDDEN,
|
|
num_hidden_layers=1,
|
|
num_target_layers=48,
|
|
):
|
|
"""Create a mocked VllmConfig for VWN model instantiation on CPU."""
|
|
hf_config = _make_hf_config(
|
|
hidden_size=hidden_size,
|
|
vwn_m=vwn_m,
|
|
vwn_r=vwn_r,
|
|
num_hidden_layers=num_hidden_layers,
|
|
)
|
|
|
|
vllm_config = MagicMock(spec=VllmConfig)
|
|
|
|
# speculative_config
|
|
vllm_config.speculative_config = MagicMock()
|
|
vllm_config.speculative_config.num_speculative_tokens = 3
|
|
vllm_config.speculative_config.draft_tensor_parallel_size = 1
|
|
vllm_config.speculative_config.parallel_drafting = False
|
|
vllm_config.speculative_config.disable_padded_drafter_batch = False
|
|
vllm_config.speculative_config.draft_model_config = MagicMock(
|
|
hf_config=hf_config,
|
|
uses_mrope=False,
|
|
uses_xdrope_dim=0,
|
|
quantization=None,
|
|
load_config=MagicMock(),
|
|
get_hidden_size=MagicMock(return_value=hidden_size),
|
|
get_inputs_embeds_size=MagicMock(return_value=hidden_size),
|
|
)
|
|
|
|
# cache_config
|
|
vllm_config.cache_config = MagicMock(spec=CacheConfig)
|
|
vllm_config.cache_config.block_size = 16
|
|
vllm_config.cache_config.kv_cache_dtype_skip_layers = None
|
|
vllm_config.cache_config.cache_dtype = "auto"
|
|
|
|
# scheduler_config
|
|
vllm_config.scheduler_config = MagicMock()
|
|
vllm_config.scheduler_config.max_num_batched_tokens = 1024
|
|
vllm_config.scheduler_config.max_num_seqs = 32
|
|
|
|
# model_config
|
|
vllm_config.model_config = MagicMock()
|
|
vllm_config.model_config.dtype = torch.float32
|
|
vllm_config.model_config.max_model_len = 2048
|
|
vllm_config.model_config.uses_mrope = False
|
|
vllm_config.model_config.uses_xdrope_dim = 0
|
|
vllm_config.model_config.enforce_eager = True
|
|
vllm_config.model_config.hf_text_config = MagicMock(spec=[])
|
|
vllm_config.model_config.hf_text_config.to_dict = MagicMock(return_value={})
|
|
vllm_config.model_config.hf_config = hf_config
|
|
vllm_config.model_config.get_num_layers = MagicMock(return_value=num_target_layers)
|
|
|
|
# compilation_config
|
|
vllm_config.compilation_config = MagicMock()
|
|
vllm_config.compilation_config.mode = CompilationMode.NONE
|
|
vllm_config.compilation_config.pass_config = MagicMock(enable_sp=False)
|
|
vllm_config.compilation_config.custom_ops = ["none"]
|
|
|
|
# parallel_config
|
|
vllm_config.parallel_config = MagicMock()
|
|
vllm_config.parallel_config.tensor_parallel_size = 1
|
|
vllm_config.parallel_config.data_parallel_rank = 0
|
|
vllm_config.parallel_config.data_parallel_size = 1
|
|
vllm_config.parallel_config.prefill_context_parallel_size = 1
|
|
vllm_config.parallel_config.decode_context_parallel_size = 1
|
|
vllm_config.parallel_config.enable_expert_parallel = False
|
|
|
|
vllm_config.additional_config = None
|
|
|
|
init_ascend_config(vllm_config)
|
|
return vllm_config
|
|
|
|
|
|
@contextmanager
|
|
def _make_model_with_mocked_ops(**kwargs):
|
|
"""Create Eagle3VwnLlamaForCausalLM with mocked attention/MLP for CPU."""
|
|
vllm_config = _create_vllm_config_for_vwn(**kwargs)
|
|
hs = vllm_config.speculative_config.draft_model_config.hf_config.hidden_size
|
|
with set_current_vllm_config(vllm_config):
|
|
model = Eagle3VwnLlamaForCausalLM(vllm_config=vllm_config, prefix="")
|
|
for layer in model.model.layers:
|
|
_mock_npu_ops_on_layer(layer)
|
|
yield model, vllm_config, hs
|
|
|
|
|
|
class TestPreVwnLayerV1:
|
|
@pytest.mark.parametrize("vwn_m,vwn_r", [(4, 1.5), (1, 1.0)])
|
|
def test_init_and_forward(self, vwn_m, vwn_r):
|
|
"""Verify layer init and forward output shape."""
|
|
vllm_config = _create_vllm_config_for_vwn(vwn_m=vwn_m, vwn_r=vwn_r)
|
|
hs, batch = _HIDDEN, 4
|
|
wd = int(hs * vwn_r)
|
|
|
|
with set_current_vllm_config(vllm_config):
|
|
layer = PreVwnLayerV1(
|
|
vllm_config=vllm_config,
|
|
prefix="test_prevwn",
|
|
config=vllm_config.speculative_config.draft_model_config.hf_config,
|
|
)
|
|
assert layer.wider_dim == wd
|
|
out = layer(torch.randn(batch, hs), torch.randn(batch, hs))
|
|
|
|
assert out.shape == (batch, wd)
|
|
|
|
|
|
class TestVwnLlamaDecoderLayer:
|
|
@pytest.mark.parametrize("vwn_m,vwn_r", [(4, 1.5), (4, 1.0)])
|
|
def test_forward_layer0(self, vwn_m, vwn_r):
|
|
"""VWN forward with various m/r configs — init + shape check."""
|
|
vllm_config = _create_vllm_config_for_vwn(vwn_m=vwn_m, vwn_r=vwn_r)
|
|
hs, batch = _HIDDEN, 4
|
|
|
|
with set_current_vllm_config(vllm_config):
|
|
layer = VwnLlamaDecoderLayer(
|
|
vllm_config=vllm_config,
|
|
prefix="model.layers.48",
|
|
config=vllm_config.speculative_config.draft_model_config.hf_config,
|
|
layer_idx=0,
|
|
)
|
|
_mock_npu_ops_on_layer(layer)
|
|
out_hidden, _ = layer(
|
|
torch.arange(batch, dtype=torch.long),
|
|
torch.randn(batch, hs),
|
|
torch.randn(batch, hs),
|
|
None,
|
|
)
|
|
|
|
assert out_hidden.shape == (batch, hs)
|
|
|
|
def test_qkv_proj_input_size_layer0(self):
|
|
"""VWN layer 0 qkv_proj input is hidden_size (not 2*hidden_size)."""
|
|
vllm_config = _create_vllm_config_for_vwn()
|
|
|
|
with set_current_vllm_config(vllm_config):
|
|
layer = VwnLlamaDecoderLayer(
|
|
vllm_config=vllm_config,
|
|
prefix="model.layers.48",
|
|
config=vllm_config.speculative_config.draft_model_config.hf_config,
|
|
layer_idx=0,
|
|
)
|
|
|
|
assert layer.self_attn.qkv_proj.input_size == _HIDDEN
|
|
|
|
|
|
class TestVwnLlamaModel:
|
|
@pytest.mark.parametrize(
|
|
"num_hidden_layers,use_input_embeds",
|
|
[
|
|
(1, False),
|
|
(1, True),
|
|
],
|
|
)
|
|
def test_forward(self, num_hidden_layers, use_input_embeds):
|
|
"""Verify layer count, type, and forward output shapes."""
|
|
vllm_config = _create_vllm_config_for_vwn(num_hidden_layers=num_hidden_layers)
|
|
hs, num_tokens = _HIDDEN, 4
|
|
|
|
with set_current_vllm_config(vllm_config):
|
|
model = VwnLlamaModel(
|
|
vllm_config=vllm_config,
|
|
prefix="model",
|
|
start_layer_id=48,
|
|
)
|
|
|
|
assert len(model.layers) == num_hidden_layers
|
|
for i, layer in enumerate(model.layers):
|
|
assert isinstance(layer, VwnLlamaDecoderLayer)
|
|
assert layer.layer_idx == i
|
|
_mock_npu_ops_on_layer(layer)
|
|
|
|
input_ids = torch.randint(0, _VOCAB, (num_tokens,))
|
|
positions = torch.arange(num_tokens, dtype=torch.long)
|
|
hidden_states = torch.randn(num_tokens, hs)
|
|
input_embeds = torch.randn(num_tokens, hs) if use_input_embeds else None
|
|
|
|
postnorm, prenorm = model(
|
|
input_ids,
|
|
positions,
|
|
hidden_states,
|
|
input_embeds=input_embeds,
|
|
)
|
|
|
|
assert postnorm.shape == (num_tokens, hs)
|
|
assert prenorm.shape == (num_tokens, hs)
|
|
|
|
|
|
class TestEagle3VwnLlamaForCausalLM:
|
|
def test_init_and_forward(self):
|
|
with _make_model_with_mocked_ops(vwn_m=4) as (model, _, hs):
|
|
assert isinstance(model.model, VwnLlamaModel)
|
|
num_tokens = 3
|
|
|
|
input_ids = torch.randint(0, _VOCAB, (num_tokens,))
|
|
positions = torch.arange(num_tokens, dtype=torch.long)
|
|
|
|
postnorm, prenorm = model(
|
|
input_ids,
|
|
positions,
|
|
torch.randn(num_tokens, hs),
|
|
)
|
|
|
|
assert postnorm.shape == (num_tokens, hs)
|
|
assert prenorm.shape == (num_tokens, hs)
|
|
|
|
def test_embed_input_ids(self):
|
|
vllm_config = _create_vllm_config_for_vwn()
|
|
num_tokens = 3
|
|
|
|
with set_current_vllm_config(vllm_config):
|
|
model = Eagle3VwnLlamaForCausalLM(vllm_config=vllm_config, prefix="")
|
|
embeds = model.embed_input_ids(torch.randint(0, _VOCAB, (num_tokens,)))
|
|
|
|
assert embeds.shape == (num_tokens, _HIDDEN)
|