tzhouamandClaude Opus 5
08884c4cfb
Four unrelated 0.29 surface changes, each reproduced locally against the pinned f5c3cc24 build: - PromptUpdateDetails.is_embed now takes only the full token list. The Bagel img2img replacement still declared (tok, seq), which crashed stage-0 engine core startup and produced the Distributed and Qwen3-Omni "no live replica" failures rather than anything KV related. - dispatch_unquantized_gemm gained a linear_backend argument and is called with it. The CPU-forcing test stubs took no arguments; they now accept and ignore whatever upstream passes so a further signature change does not break them again. - VocabParallelEmbedding gained a fused kernel that engages whenever tp_size > 1 on CUDA, and it has no CPU implementation. The encoder TP fixtures drive CPU tensors under a faked TP group, so they pin the portable masked path. - The pyav video decoder backend was removed; valid options are now opencv, torchcodec, pynvvideocodec and deepstream. The OpenAI video endpoint hardcoded pyav, which turned every video request into a 400. It moves to opencv, the upstream default. 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>
556 lines
21 KiB
Python
556 lines
21 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import torch
|
|
from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config
|
|
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
|
|
|
|
from vllm_omni.diffusion.models.mistral_encoder.mistral_encoder import (
|
|
MistralEncoderModel,
|
|
MistralEncoderOutput,
|
|
MistralRotaryEmbedding,
|
|
)
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
_MODULE = "vllm_omni.diffusion.models.mistral_encoder.mistral_encoder"
|
|
|
|
SMALL_MISTRAL_CONFIG = dict(
|
|
hidden_size=64,
|
|
num_attention_heads=8,
|
|
num_key_value_heads=4,
|
|
head_dim=8,
|
|
intermediate_size=128,
|
|
num_hidden_layers=2,
|
|
rms_norm_eps=1e-5,
|
|
max_position_embeddings=512,
|
|
rope_theta=1000000.0,
|
|
vocab_size=256,
|
|
)
|
|
|
|
|
|
def _make_config(**overrides):
|
|
return SimpleNamespace(**{**SMALL_MISTRAL_CONFIG, **overrides})
|
|
|
|
|
|
def _make_nested_config(**overrides):
|
|
"""Simulate Mistral3Config with a text_config attribute."""
|
|
text_config = _make_config(**overrides)
|
|
return SimpleNamespace(text_config=text_config)
|
|
|
|
|
|
@pytest.fixture(scope="function", autouse=True)
|
|
def setup_tp_group(monkeypatch, mocker):
|
|
"""Set up TP=2, rank=0, and VllmConfig for all tests."""
|
|
device_config = DeviceConfig(device="cpu")
|
|
|
|
monkeypatch.setattr(
|
|
"vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size",
|
|
lambda: 2,
|
|
)
|
|
monkeypatch.setattr(
|
|
"vllm.model_executor.layers.linear.get_tensor_model_parallel_rank",
|
|
lambda: 0,
|
|
)
|
|
monkeypatch.setattr(f"{_MODULE}.get_tensor_model_parallel_world_size", lambda: 2)
|
|
monkeypatch.setattr(
|
|
"vllm.model_executor.layers.vocab_parallel_embedding.get_tensor_model_parallel_world_size",
|
|
lambda: 2,
|
|
)
|
|
monkeypatch.setattr(
|
|
"vllm.model_executor.layers.vocab_parallel_embedding.get_tensor_model_parallel_rank",
|
|
lambda: 0,
|
|
)
|
|
|
|
# vLLM 0.29 added a CUDA-only fused embedding kernel that engages whenever
|
|
# tp_size > 1; these tests drive CPU tensors, so keep the portable masked path.
|
|
monkeypatch.setattr(
|
|
"vllm.model_executor.layers.vocab_parallel_embedding.current_platform.is_cuda",
|
|
lambda: False,
|
|
)
|
|
|
|
mock_tp_group = mocker.MagicMock()
|
|
mock_tp_group.world_size = 2
|
|
mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_tp_group)
|
|
|
|
# Mock TP communication ops used during forward passes. Each op is
|
|
# imported by reference into the modules that use it, so we must patch
|
|
# at every import site.
|
|
_identity = lambda x: x # noqa: E731
|
|
monkeypatch.setattr(
|
|
"vllm.model_executor.layers.vocab_parallel_embedding.tensor_model_parallel_all_reduce",
|
|
_identity,
|
|
)
|
|
monkeypatch.setattr(
|
|
"vllm.model_executor.layers.linear.tensor_model_parallel_all_reduce",
|
|
_identity,
|
|
)
|
|
monkeypatch.setattr(
|
|
"vllm.model_executor.layers.linear.tensor_model_parallel_all_gather",
|
|
_identity,
|
|
)
|
|
monkeypatch.setattr(
|
|
f"{_MODULE}.tensor_model_parallel_all_gather",
|
|
_identity,
|
|
)
|
|
mocker.patch("torch.distributed.broadcast")
|
|
|
|
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,
|
|
)
|
|
|
|
with set_current_vllm_config(VllmConfig(device_config=device_config)):
|
|
yield
|
|
|
|
|
|
class TestConfigParsing:
|
|
"""Verify that MistralEncoderModel extracts config correctly."""
|
|
|
|
def test_plain_config(self):
|
|
config = _make_config()
|
|
model = MistralEncoderModel(config, prefix="text_encoder")
|
|
|
|
assert model.hidden_size == 64
|
|
assert model.num_heads == 8
|
|
assert model.num_kv_heads == 4
|
|
assert model.head_dim == 8
|
|
assert model.intermediate_size == 128
|
|
assert model.num_layers == 2
|
|
assert model.rms_norm_eps == 1e-5
|
|
assert model.rope_theta == 1000000.0
|
|
assert model.vocab_size == 256
|
|
|
|
def test_nested_text_config(self):
|
|
config = _make_nested_config()
|
|
model = MistralEncoderModel(config, prefix="text_encoder")
|
|
|
|
assert model.hidden_size == 64
|
|
assert model.num_heads == 8
|
|
assert model.num_kv_heads == 4
|
|
assert model.config is config.text_config
|
|
|
|
def test_defaults_when_fields_missing(self):
|
|
config = SimpleNamespace(
|
|
hidden_size=64,
|
|
num_attention_heads=8,
|
|
intermediate_size=128,
|
|
num_hidden_layers=1,
|
|
vocab_size=256,
|
|
)
|
|
model = MistralEncoderModel(config, prefix="text_encoder")
|
|
|
|
assert model.num_kv_heads == 8, "should fall back to num_attention_heads"
|
|
assert model.head_dim == 8, "should compute hidden_size // num_heads"
|
|
assert model.rms_norm_eps == 1e-5
|
|
assert model.max_position_embeddings == 131072
|
|
assert model.rope_theta == 1000000.0
|
|
|
|
|
|
class TestRoPEInitialization:
|
|
"""Verify that RoPE inv_freq is computed from config, not left uninitialized."""
|
|
|
|
def test_inv_freq_deterministic(self):
|
|
rope = MistralRotaryEmbedding(head_dim=8, max_position_embeddings=512, rope_theta=1000000.0)
|
|
|
|
expected = 1.0 / (1000000.0 ** (torch.arange(0, 8, 2, dtype=torch.float32) / 8))
|
|
assert torch.allclose(rope.inv_freq, expected)
|
|
|
|
def test_different_theta_produces_different_freqs(self):
|
|
rope_a = MistralRotaryEmbedding(head_dim=8, max_position_embeddings=512, rope_theta=10000.0)
|
|
rope_b = MistralRotaryEmbedding(head_dim=8, max_position_embeddings=512, rope_theta=1000000.0)
|
|
|
|
assert not torch.allclose(rope_a.inv_freq, rope_b.inv_freq)
|
|
|
|
def test_cos_sin_shape_and_identity(self):
|
|
rope = MistralRotaryEmbedding(head_dim=8, max_position_embeddings=512, rope_theta=1000000.0)
|
|
seq_len = 16
|
|
position_ids = torch.arange(seq_len).unsqueeze(0)
|
|
cos, sin = rope(position_ids, torch.float32)
|
|
|
|
assert cos.shape == (1, seq_len, 8)
|
|
assert sin.shape == (1, seq_len, 8)
|
|
assert torch.allclose(cos**2 + sin**2, torch.ones_like(cos), atol=1e-6)
|
|
|
|
def test_model_rope_uses_config_theta(self):
|
|
config = _make_config(rope_theta=10000.0)
|
|
model = MistralEncoderModel(config, prefix="text_encoder")
|
|
|
|
expected = 1.0 / (10000.0 ** (torch.arange(0, 8, 2, dtype=torch.float32) / 8))
|
|
actual = model.language_model.model.rotary_emb.inv_freq
|
|
assert torch.allclose(actual, expected)
|
|
|
|
|
|
class TestWeightLoading:
|
|
"""Test weight loading and stacked params mapping."""
|
|
|
|
def test_qkv_weights_loaded(self):
|
|
config = _make_config(num_hidden_layers=1)
|
|
model = MistralEncoderModel(config, prefix="text_encoder")
|
|
|
|
hidden_size = config.hidden_size
|
|
num_heads = config.num_attention_heads
|
|
num_kv_heads = config.num_key_value_heads
|
|
head_dim = config.head_dim
|
|
|
|
prefix = "language_model.model.layers.0.self_attn."
|
|
weights = [
|
|
(prefix + "q_proj.weight", torch.randn(num_heads * head_dim, hidden_size)),
|
|
(prefix + "k_proj.weight", torch.randn(num_kv_heads * head_dim, hidden_size)),
|
|
(prefix + "v_proj.weight", torch.randn(num_kv_heads * head_dim, hidden_size)),
|
|
]
|
|
|
|
loaded = model.load_weights(weights)
|
|
assert len(loaded) > 0
|
|
assert any("qkv_proj" in p for p in loaded)
|
|
|
|
attn = model.language_model.model.layers[0].self_attn
|
|
# TP=2: q sharded to num_heads/2, kv sharded to num_kv_heads/2
|
|
expected_dim = (num_heads // 2 + 2 * (num_kv_heads // 2)) * head_dim
|
|
assert attn.qkv_proj.weight.shape == (expected_dim, hidden_size)
|
|
|
|
def test_gate_up_weights_loaded(self):
|
|
config = _make_config(num_hidden_layers=1)
|
|
model = MistralEncoderModel(config, prefix="text_encoder")
|
|
|
|
hidden_size = config.hidden_size
|
|
intermediate_size = config.intermediate_size
|
|
|
|
prefix = "language_model.model.layers.0.mlp."
|
|
weights = [
|
|
(prefix + "gate_proj.weight", torch.randn(intermediate_size, hidden_size)),
|
|
(prefix + "up_proj.weight", torch.randn(intermediate_size, hidden_size)),
|
|
]
|
|
|
|
loaded = model.load_weights(weights)
|
|
assert len(loaded) > 0
|
|
assert any("gate_up_proj" in p for p in loaded)
|
|
|
|
mlp = model.language_model.model.layers[0].mlp
|
|
# TP=2: each shard is intermediate_size/2, two merged
|
|
expected_dim = 2 * (intermediate_size // 2)
|
|
assert mlp.gate_up_proj.weight.shape == (expected_dim, hidden_size)
|
|
|
|
def test_skips_vision_but_loads_lm_head(self):
|
|
config = _make_config(num_hidden_layers=1)
|
|
model = MistralEncoderModel(config, prefix="text_encoder")
|
|
|
|
weights = [
|
|
("language_model.lm_head.weight", torch.randn(256, 64)),
|
|
("vision_tower.encoder.weight", torch.randn(64, 64)),
|
|
("multi_modal_projector.linear.weight", torch.randn(64, 64)),
|
|
]
|
|
|
|
loaded = model.load_weights(weights)
|
|
assert any("lm_head" in p for p in loaded)
|
|
assert not any("vision_tower" in p for p in loaded)
|
|
assert not any("multi_modal_projector" in p for p in loaded)
|
|
|
|
def test_unknown_weights_ignored(self):
|
|
config = _make_config(num_hidden_layers=1)
|
|
model = MistralEncoderModel(config, prefix="text_encoder")
|
|
|
|
weights = [("totally.fake.weight", torch.randn(10, 10))]
|
|
loaded = model.load_weights(weights)
|
|
assert len(loaded) == 0
|
|
|
|
|
|
class TestModelStructure:
|
|
"""Verify module hierarchy matches HF checkpoint layout."""
|
|
|
|
def test_module_nesting(self):
|
|
config = _make_config(num_hidden_layers=2)
|
|
model = MistralEncoderModel(config, prefix="text_encoder")
|
|
|
|
assert hasattr(model, "language_model")
|
|
assert hasattr(model.language_model, "model")
|
|
m = model.language_model.model
|
|
assert hasattr(m, "embed_tokens")
|
|
assert hasattr(m, "layers")
|
|
assert hasattr(m, "norm")
|
|
assert hasattr(m, "rotary_emb")
|
|
assert len(m.layers) == 2
|
|
|
|
def test_param_names_match_checkpoint_prefix(self):
|
|
config = _make_config(num_hidden_layers=1)
|
|
model = MistralEncoderModel(config, prefix="text_encoder")
|
|
|
|
param_names = set(dict(model.named_parameters()).keys())
|
|
assert any(n.startswith("language_model.model.embed_tokens") for n in param_names)
|
|
assert any(n.startswith("language_model.model.layers.0.self_attn") for n in param_names)
|
|
assert any(n.startswith("language_model.model.layers.0.mlp") for n in param_names)
|
|
|
|
|
|
class TestKVCache:
|
|
"""Verify KV cache plumbing through attention, layer, and model."""
|
|
|
|
def test_attention_returns_none_kv_when_cache_off(self):
|
|
from vllm_omni.diffusion.models.mistral_encoder.mistral_encoder import (
|
|
MistralEncoderAttention,
|
|
)
|
|
|
|
attn = MistralEncoderAttention(
|
|
hidden_size=64,
|
|
num_heads=8,
|
|
num_kv_heads=4,
|
|
head_dim=8,
|
|
prefix="test",
|
|
)
|
|
hidden = torch.randn(1, 4, 64)
|
|
cos = torch.ones(1, 4, 8)
|
|
sin = torch.zeros(1, 4, 8)
|
|
out, kv = attn(hidden, cos, sin, use_cache=False)
|
|
assert kv is None
|
|
assert out.shape == (1, 4, 64)
|
|
|
|
def test_attention_returns_kv_when_cache_on(self):
|
|
from vllm_omni.diffusion.models.mistral_encoder.mistral_encoder import (
|
|
MistralEncoderAttention,
|
|
)
|
|
|
|
attn = MistralEncoderAttention(
|
|
hidden_size=64,
|
|
num_heads=8,
|
|
num_kv_heads=4,
|
|
head_dim=8,
|
|
prefix="test",
|
|
)
|
|
hidden = torch.randn(1, 4, 64)
|
|
cos = torch.ones(1, 4, 8)
|
|
sin = torch.zeros(1, 4, 8)
|
|
out, kv = attn(hidden, cos, sin, use_cache=True)
|
|
assert kv is not None
|
|
k, v = kv
|
|
# TP=2: num_kv_heads shard = 4//2 = 2
|
|
assert k.shape == (1, 2, 4, 8)
|
|
assert v.shape == (1, 2, 4, 8)
|
|
|
|
def test_attention_appends_past_kv(self):
|
|
from vllm_omni.diffusion.models.mistral_encoder.mistral_encoder import (
|
|
MistralEncoderAttention,
|
|
)
|
|
|
|
attn = MistralEncoderAttention(
|
|
hidden_size=64,
|
|
num_heads=8,
|
|
num_kv_heads=4,
|
|
head_dim=8,
|
|
prefix="test",
|
|
)
|
|
# Simulate a prefill with 4 tokens
|
|
hidden = torch.randn(1, 4, 64)
|
|
cos = torch.ones(1, 4, 8)
|
|
sin = torch.zeros(1, 4, 8)
|
|
_, kv = attn(hidden, cos, sin, use_cache=True)
|
|
|
|
# Simulate a decode step with 1 token
|
|
hidden_new = torch.randn(1, 1, 64)
|
|
cos_new = torch.ones(1, 1, 8)
|
|
sin_new = torch.zeros(1, 1, 8)
|
|
out, kv2 = attn(hidden_new, cos_new, sin_new, past_key_value=kv, use_cache=True)
|
|
k2, v2 = kv2
|
|
assert k2.shape == (1, 2, 5, 8), "should be past(4) + new(1)"
|
|
assert out.shape == (1, 1, 64)
|
|
|
|
def test_model_forward_use_cache(self):
|
|
config = _make_config(num_hidden_layers=2)
|
|
model = MistralEncoderModel(config, prefix="text_encoder")
|
|
input_ids = torch.randint(0, 128, (1, 8))
|
|
|
|
output = model(input_ids, use_cache=True)
|
|
assert isinstance(output, MistralEncoderOutput)
|
|
assert output.past_key_values is not None
|
|
assert len(output.past_key_values) == 2
|
|
# Each layer's cache: (k, v) with seq_len=8
|
|
k, v = output.past_key_values[0]
|
|
assert k.shape[2] == 8
|
|
assert v.shape[2] == 8
|
|
|
|
def test_model_forward_no_cache(self):
|
|
config = _make_config(num_hidden_layers=1)
|
|
model = MistralEncoderModel(config, prefix="text_encoder")
|
|
input_ids = torch.randint(0, 128, (1, 4))
|
|
|
|
output = model(input_ids, use_cache=False)
|
|
assert output.past_key_values is None
|
|
|
|
def test_model_decode_with_past(self):
|
|
config = _make_config(num_hidden_layers=1)
|
|
model = MistralEncoderModel(config, prefix="text_encoder")
|
|
|
|
# Prefill
|
|
input_ids = torch.randint(0, 128, (1, 4))
|
|
output = model(input_ids, use_cache=True)
|
|
past = output.past_key_values
|
|
|
|
# Decode one token
|
|
new_token = torch.randint(0, 128, (1, 1))
|
|
output2 = model(new_token, use_cache=True, past_key_values=past)
|
|
assert output2.last_hidden_state.shape == (1, 1, 64)
|
|
k2, v2 = output2.past_key_values[0]
|
|
assert k2.shape[2] == 5, "cache should grow from 4 to 5"
|
|
|
|
|
|
class TestRoPEOffset:
|
|
"""Verify RoPE offset produces correct positions for decode steps."""
|
|
|
|
def test_offset_zero_matches_original(self):
|
|
rope = MistralRotaryEmbedding(head_dim=8, max_position_embeddings=512, rope_theta=1000000.0)
|
|
position_ids = torch.arange(4).unsqueeze(0)
|
|
cos_a, sin_a = rope(position_ids, torch.float32)
|
|
cos_b, sin_b = rope(position_ids, torch.float32)
|
|
assert torch.allclose(cos_a, cos_b)
|
|
assert torch.allclose(sin_a, sin_b)
|
|
|
|
def test_offset_produces_shifted_positions(self):
|
|
rope = MistralRotaryEmbedding(head_dim=8, max_position_embeddings=512, rope_theta=1000000.0)
|
|
# Full sequence of 5 positions
|
|
full_ids = torch.arange(5).unsqueeze(0)
|
|
cos_full, sin_full = rope(full_ids, torch.float32)
|
|
# Position 4 only
|
|
offset_ids = torch.tensor([[4]])
|
|
cos_off, sin_off = rope(offset_ids, torch.float32)
|
|
assert torch.allclose(cos_full[:, 4:5], cos_off)
|
|
assert torch.allclose(sin_full[:, 4:5], sin_off)
|
|
|
|
|
|
class TestGenerate:
|
|
"""Verify autoregressive generate() method."""
|
|
|
|
def test_generate_produces_tokens(self):
|
|
config = _make_config(num_hidden_layers=1)
|
|
model = MistralEncoderModel(config, prefix="text_encoder")
|
|
input_ids = torch.randint(0, 128, (1, 4))
|
|
attention_mask = torch.ones(1, 4, dtype=torch.long)
|
|
|
|
output = model.generate(
|
|
input_ids=input_ids,
|
|
attention_mask=attention_mask,
|
|
max_new_tokens=3,
|
|
do_sample=False,
|
|
use_cache=True,
|
|
)
|
|
assert output.shape[0] == 1
|
|
# prompt(4) + generated(3) = 7, unless EOS hit early
|
|
assert output.shape[1] >= 5
|
|
assert output.shape[1] <= 7
|
|
# prompt tokens preserved
|
|
assert torch.equal(output[:, :4], input_ids)
|
|
|
|
def test_generate_stops_at_eos(self):
|
|
config = _make_config(num_hidden_layers=1, eos_token_id=0)
|
|
model = MistralEncoderModel(config, prefix="text_encoder")
|
|
|
|
# Manually set embed_tokens so that token 0 maps to a specific embedding
|
|
# that will produce logits strongly favouring token 0 again.
|
|
# This is a probabilistic test — we just verify it doesn't exceed max.
|
|
input_ids = torch.randint(1, 128, (1, 4))
|
|
output = model.generate(
|
|
input_ids=input_ids,
|
|
max_new_tokens=10,
|
|
do_sample=False,
|
|
eos_token_id=0,
|
|
)
|
|
# Should have at most prompt(4) + max_new(10) = 14 tokens
|
|
assert output.shape[1] <= 14
|
|
|
|
def test_generate_greedy_deterministic(self):
|
|
config = _make_config(num_hidden_layers=1)
|
|
model = MistralEncoderModel(config, prefix="text_encoder")
|
|
input_ids = torch.randint(0, 128, (1, 4))
|
|
|
|
out1 = model.generate(input_ids=input_ids, max_new_tokens=5, do_sample=False)
|
|
out2 = model.generate(input_ids=input_ids, max_new_tokens=5, do_sample=False)
|
|
assert torch.equal(out1, out2)
|
|
|
|
def test_generate_ignores_extra_kwargs(self):
|
|
"""generate() should accept and ignore pixel_values and other HF kwargs."""
|
|
config = _make_config(num_hidden_layers=1)
|
|
model = MistralEncoderModel(config, prefix="text_encoder")
|
|
input_ids = torch.randint(0, 128, (1, 4))
|
|
|
|
output = model.generate(
|
|
input_ids=input_ids,
|
|
max_new_tokens=2,
|
|
do_sample=False,
|
|
pixel_values=torch.randn(1, 3, 224, 224),
|
|
)
|
|
assert output.shape[1] >= 5
|
|
|
|
|
|
class TestComputeLogits:
|
|
"""Verify logits computation via lm_head weight."""
|
|
|
|
def test_logits_shape(self):
|
|
config = _make_config(num_hidden_layers=1, vocab_size=256)
|
|
model = MistralEncoderModel(config, prefix="text_encoder")
|
|
|
|
hidden = torch.randn(1, 4, 64)
|
|
logits = model._compute_logits(hidden)
|
|
# TP=2: VocabParallelEmbedding stores vocab_size/2 = 128 per shard
|
|
# With TP=2 but mocked (no actual all_gather), local logits only
|
|
# In real TP, all_gather would give (1, 4, 256)
|
|
# With mock, tp_size=2 triggers all_gather path but the mock may not
|
|
# actually gather. Just verify we get a tensor back.
|
|
assert logits.dim() == 3
|
|
assert logits.shape[0] == 1
|
|
assert logits.shape[1] == 4
|
|
|
|
|
|
class TestQuantConfig:
|
|
"""Verify quant_config is threaded into weight linears only, with correct prefixes."""
|
|
|
|
def test_quant_config_on_weight_linears(self, mocker):
|
|
quant_config = mocker.MagicMock(name="QuantizationConfig")
|
|
# Keep construction on the CPU-safe unquantized path while still
|
|
# verifying LinearBase consults the supplied quant_config.
|
|
quant_config.get_quant_method.return_value = UnquantizedLinearMethod()
|
|
|
|
config = _make_config(num_hidden_layers=1)
|
|
model = MistralEncoderModel(config, prefix="text_encoder", quant_config=quant_config)
|
|
|
|
layer = model.language_model.model.layers[0]
|
|
for linear in (
|
|
layer.self_attn.qkv_proj,
|
|
layer.self_attn.o_proj,
|
|
layer.mlp.gate_up_proj,
|
|
layer.mlp.down_proj,
|
|
):
|
|
assert linear.quant_config is quant_config
|
|
|
|
# get_quant_method should be consulted for the four weight linears.
|
|
assert quant_config.get_quant_method.call_count == 4
|
|
seen_prefixes = set()
|
|
for call in quant_config.get_quant_method.call_args_list:
|
|
if "prefix" in call.kwargs:
|
|
seen_prefixes.add(call.kwargs["prefix"])
|
|
elif len(call.args) > 1:
|
|
seen_prefixes.add(call.args[1])
|
|
assert any(p.endswith("self_attn.qkv_proj") for p in seen_prefixes)
|
|
assert any(p.endswith("self_attn.o_proj") for p in seen_prefixes)
|
|
assert any(p.endswith("mlp.gate_up_proj") for p in seen_prefixes)
|
|
assert any(p.endswith("mlp.down_proj") for p in seen_prefixes)
|
|
|
|
# Embeddings / lm_head stay unquantized (quant_config never passed).
|
|
embed = model.language_model.model.embed_tokens
|
|
lm_head = model.language_model.lm_head
|
|
assert getattr(embed, "quant_config", None) is None
|
|
assert getattr(lm_head, "quant_config", None) is None
|
|
|
|
def test_layer_prefixes_include_text_encoder(self):
|
|
config = _make_config(num_hidden_layers=2)
|
|
model = MistralEncoderModel(config, prefix="text_encoder")
|
|
|
|
layer0 = model.language_model.model.layers[0]
|
|
assert layer0.self_attn.qkv_proj.prefix.startswith("text_encoder.language_model.model.layers.0")
|
|
assert layer0.self_attn.o_proj.prefix.startswith("text_encoder.language_model.model.layers.0")
|
|
assert layer0.mlp.gate_up_proj.prefix.startswith("text_encoder.language_model.model.layers.0")
|
|
assert layer0.mlp.down_proj.prefix.startswith("text_encoder.language_model.model.layers.0")
|
|
|
|
layer1 = model.language_model.model.layers[1]
|
|
assert "layers.1" in layer1.self_attn.qkv_proj.prefix
|