[Bugfix] Fix Dense module loading for sentence-transformers embedding models v3

Signed-off-by: FFFfff1FFFfff <[email protected]>

Author: FFFfff1FFFfff

tests/model_executor/test_st_projector_unit.py

@@ -0,0 +1,156 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+from unittest.mock import Mock, patch
+from vllm.model_executor.layers.pooler import EmbeddingPoolerHead
+
+
+def test_embedding_pooler_head_with_projector():
+    """Test EmbeddingPoolerHead with projector."""
+    # Create a mock projector
+    projector = torch.nn.Sequential(
+        torch.nn.Linear(1024, 1792),
+        torch.nn.GELU()
+    )
+    
+    # Create EmbeddingPoolerHead with projector
+    head = EmbeddingPoolerHead(projector=projector)
+    
+    # Test with single tensor
+    test_input = torch.randn(2, 1024)
+    
+    # Mock the pooling parameters
+    class MockPoolingParam:
+        def __init__(self):
+            self.dimensions = None
+            self.normalize = True
+    
+    # Mock metadata
+    class MockMetadata:
+        def __init__(self):
+            pass
+    
+    mock_metadata = MockMetadata()
+    
+    # Mock get_pooling_params function
+    with patch('vllm.model_executor.layers.pooler.get_pooling_params') as mock_get_params:
+        mock_get_params.return_value = [MockPoolingParam()]
+        
+        # Call forward - this will test the projector path
+        output = head(test_input, mock_metadata)
+        
+        # Should apply projector first (1024 -> 1792), then normalize
+        assert output.shape == (2, 1792), f"Expected (2, 1792), got {output.shape}"
+        
+        # Check that output is normalized (L2 norm should be ~1)
+        norms = torch.norm(output, p=2, dim=-1)
+        assert torch.allclose(norms, torch.ones_like(norms), atol=1e-5), \
+            "Output should be normalized"
+
+
+def test_embedding_pooler_head_without_projector():
+    """Test EmbeddingPoolerHead without projector (regression test)."""
+    # Create EmbeddingPoolerHead without projector
+    head = EmbeddingPoolerHead(projector=None)
+    
+    # Test with single tensor
+    test_input = torch.randn(2, 1024)
+    
+    # Mock the pooling parameters
+    class MockPoolingParam:
+        def __init__(self):
+            self.dimensions = None
+            self.normalize = True
+    
+    # Mock metadata
+    class MockMetadata:
+        def __init__(self):
+            pass
+    
+    mock_metadata = MockMetadata()
+    
+    # Mock get_pooling_params function
+    with patch('vllm.model_executor.layers.pooler.get_pooling_params') as mock_get_params:
+        mock_get_params.return_value = [MockPoolingParam()]
+        
+        # Call forward - this should just normalize without projection
+        output = head(test_input, mock_metadata)
+        
+        # Should maintain the input dimension
+        assert output.shape == (2, 1024), f"Expected (2, 1024), got {output.shape}"
+        
+        # Check that output is normalized
+        norms = torch.norm(output, p=2, dim=-1)
+        assert torch.allclose(norms, torch.ones_like(norms), atol=1e-5), \
+            "Output should be normalized"
+
+
+def test_embedding_pooler_head_dimension_mismatch():
+    """Test that dimension mismatch raises an error."""
+    # Create a projector expecting 512-dim input
+    projector = torch.nn.Sequential(
+        torch.nn.Linear(512, 1792),  # Expects 512, not 1024
+        torch.nn.GELU()
+    )
+    
+    head = EmbeddingPoolerHead(projector=projector)
+    
+    # Test with wrong dimension (1024 instead of 512)
+    test_input = torch.randn(2, 1024)
+    
+    class MockPoolingParam:
+        def __init__(self):
+            self.dimensions = None
+            self.normalize = True
+    
+    class MockMetadata:
+        def __init__(self):
+            pass
+    
+    mock_metadata = MockMetadata()
+    
+    with patch('vllm.model_executor.layers.pooler.get_pooling_params') as mock_get_params:
+        mock_get_params.return_value = [MockPoolingParam()]
+        
+        # Should raise ValueError due to dimension mismatch
+        with pytest.raises(ValueError, match="Dimension mismatch"):
+            head(test_input, mock_metadata)
+
+
+def test_embedding_pooler_head_device_handling():
+    """Test that projector moves to correct device."""
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA not available")
+    
+    # Create projector on CPU
+    projector = torch.nn.Sequential(
+        torch.nn.Linear(1024, 1792),
+        torch.nn.GELU()
+    )
+    
+    head = EmbeddingPoolerHead(projector=projector)
+    
+    # Test input on CUDA
+    test_input = torch.randn(2, 1024).cuda()
+    
+    class MockPoolingParam:
+        def __init__(self):
+            self.dimensions = None
+            self.normalize = True
+    
+    class MockMetadata:
+        def __init__(self):
+            pass
+    
+    mock_metadata = MockMetadata()
+    
+    with patch('vllm.model_executor.layers.pooler.get_pooling_params') as mock_get_params:
+        mock_get_params.return_value = [MockPoolingParam()]
+        
+        output = head(test_input, mock_metadata)
+        
+        # Output should be on CUDA
+        assert output.is_cuda, "Output should be on CUDA"
+        assert output.shape == (2, 1792)

tests/models/language/pooling/test_st_projector.py

@@ -0,0 +1,218 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from typing import Any
+import pytest
+import numpy as np
+from scipy.spatial.distance import cosine
+from ...utils import EmbedModelInfo
+from .mteb_utils import mteb_test_embed_models
+
+
+def _get_vllm_embeddings(vllm_runner, model_info: EmbedModelInfo, test_texts: list[str]):
+    """Helper function to get vLLM embeddings."""
+    vllm_extra_kwargs: dict[str, Any] = {}
+    if model_info.architecture == "GteNewModel":
+        vllm_extra_kwargs["hf_overrides"] = {"architectures": ["GteNewModel"]}
+        
+    with vllm_runner(model_info.name,
+                     runner="pooling",
+                     max_model_len=None,
+                     trust_remote_code=True,
+                     **vllm_extra_kwargs) as vllm_model:
+        embeddings = vllm_model.encode(test_texts)
+        
+        # Extract embedding data
+        data = []
+        for emb in embeddings:
+            if hasattr(emb, 'outputs'):
+                data.append(emb.outputs.data.cpu().numpy())
+            else:
+                data.append(emb.cpu().numpy() if hasattr(emb, 'cpu') else emb)
+        return np.array(data)
+
+
+def _get_hf_embeddings(hf_runner, model_info: EmbedModelInfo, test_texts: list[str]):
+    """Helper function to get HuggingFace embeddings."""
+    with hf_runner(model_info.name,
+                   is_sentence_transformer=True,
+                   dtype="float32") as hf_model:
+        embeddings = hf_model.encode(test_texts)
+        if hasattr(embeddings, 'cpu'):
+            return embeddings.cpu().numpy()
+        else:
+            return np.array(embeddings)
+
+
+# Test models with ST projectors (Dense layers)
+ST_PROJECTOR_MODELS = [
+    EmbedModelInfo("TencentBAC/Conan-embedding-v1",
+                architecture="BertModel",
+                enable_test=True),
+    # Add more ST models with projectors as they become available
+]
+
+# Test models without ST projectors (for regression testing)
+NON_PROJECTOR_MODELS = [
+    EmbedModelInfo("thenlper/gte-large",
+                   architecture="BertModel",
+                   enable_test=True),
+    EmbedModelInfo("Alibaba-NLP/gte-base-en-v1.5",
+                   architecture="GteNewModel",
+                   enable_test=True),
+    EmbedModelInfo("Qwen/Qwen3-Embedding-0.6B",
+                   architecture="Qwen3ForCausalLM",
+                   dtype="float32",
+                   enable_test=True),
+]
+
+@pytest.mark.parametrize("model_info", ST_PROJECTOR_MODELS)
+def test_st_projector_loading(vllm_runner, model_info: EmbedModelInfo) -> None:
+    """Test that ST projector models load correctly with their projectors."""
+    if not model_info.enable_test:
+        pytest.skip("Skipping test.")
+    
+    test_texts = ["This is a test sentence."]
+    embeddings_data = _get_vllm_embeddings(vllm_runner, model_info, test_texts)
+    
+    # Check if dimensions match expected projector output
+    actual_dim = embeddings_data.shape[-1]
+    expected_dim = 1792
+    assert actual_dim == expected_dim, f"Expected {expected_dim}, got {actual_dim}"
+
+@pytest.mark.parametrize("model_info", ST_PROJECTOR_MODELS)  
+def test_compare_with_hf_dimensions(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:
+    """Compare embedding dimensions between vLLM and HuggingFace."""
+    if not model_info.enable_test:
+        pytest.skip("Skipping test.")
+    
+    test_texts = ["This is a test sentence for dimension comparison."]
+    
+    # Get embeddings from both implementations
+    vllm_data = _get_vllm_embeddings(vllm_runner, model_info, test_texts)
+    hf_data = _get_hf_embeddings(hf_runner, model_info, test_texts)
+    
+    # Compare dimensions
+    vllm_dim = vllm_data.shape[-1]
+    hf_dim = hf_data.shape[-1]
+    
+    assert vllm_dim == hf_dim, f"Embedding dimension mismatch: vLLM {vllm_dim} vs HF {hf_dim}"
+    print(f"✓ Embedding dimensions match: {vllm_dim}")
+
+
+@pytest.mark.parametrize("model_info", ST_PROJECTOR_MODELS)
+def test_embedding_numerical_similarity(hf_runner, vllm_runner, 
+                                       model_info: EmbedModelInfo) -> None:
+    """Test numerical similarity between vLLM and HuggingFace embeddings."""
+    if not model_info.enable_test:
+        pytest.skip("Skipping test.")
+    
+    test_texts = [
+        "This is a test sentence for numerical comparison.",
+        "Another sentence to verify embedding quality.",
+        "机器学习是人工智能的一个重要分支。",  # Chinese test
+    ]
+    
+    # Get embeddings from both implementations
+    vllm_data = _get_vllm_embeddings(vllm_runner, model_info, test_texts)
+    hf_data = _get_hf_embeddings(hf_runner, model_info, test_texts)
+    
+    # Verify shape matching
+    assert vllm_data.shape == hf_data.shape, f"Shape mismatch: vLLM {vllm_data.shape} vs HF {hf_data.shape}"
+    
+    print(f"Embedding shape: {vllm_data.shape}")
+    print(f"Embedding dimension: {vllm_data.shape[-1]}")
+    
+    # Compute similarities for each test text
+    similarities = []
+    for i, text in enumerate(test_texts):
+        vllm_emb = vllm_data[i]
+        hf_emb = hf_data[i]
+        
+        # Compute cosine similarity
+        similarity = 1 - cosine(vllm_emb, hf_emb)
+        similarities.append(similarity)
+        
+        print(f"Text {i+1}: '{text[:50]}{'...' if len(text) > 50 else ''}'")
+        print(f"  Cosine similarity: {similarity:.6f}")
+        
+        # Verify similarity threshold
+        min_similarity = 0.95
+        assert similarity > min_similarity, (
+            f"Text {i+1} similarity too low: {similarity:.6f} < {min_similarity}\n"
+            f"vLLM norm: {np.linalg.norm(vllm_emb):.6f}, "
+            f"HF norm: {np.linalg.norm(hf_emb):.6f}")
+    
+    # Verify average similarity
+    avg_similarity = np.mean(similarities)
+    print(f"\nAverage cosine similarity: {avg_similarity:.6f}")
+    
+    assert avg_similarity > 0.98, (
+        f"Average similarity too low: {avg_similarity:.6f} < 0.98")
+    
+    print("✓ All numerical similarity tests passed!")
+
+
+@pytest.mark.parametrize("model_info", ST_PROJECTOR_MODELS)
+def test_embedding_quality_checks(vllm_runner, model_info: EmbedModelInfo) -> None:
+    """Test embedding quality: non-zero, non-constant, and distinct vectors."""
+    if not model_info.enable_test:
+        pytest.skip("Skipping test.")
+    
+    test_texts = [
+        "First test sentence.",
+        "Second different sentence.",
+        "Completely different content here."
+    ]
+    
+    embeddings_data = _get_vllm_embeddings(vllm_runner, model_info, test_texts)
+    
+    print(f"Embeddings shape: {embeddings_data.shape}")
+    
+    # Verify non-zero vectors
+    for i, emb in enumerate(embeddings_data):
+        norm = np.linalg.norm(emb)
+        print(f"Embedding {i+1} L2 norm: {norm:.6f}")
+        assert norm > 1e-6, f"Embedding {i+1} is too close to zero vector: norm={norm}"
+        
+        # Verify non-constant vectors
+        std = np.std(emb)
+        print(f"Embedding {i+1} std: {std:.6f}")
+        assert std > 1e-6, f"Embedding {i+1} is too close to constant vector: std={std}"
+    
+    # Verify different texts produce different embeddings
+    for i in range(len(embeddings_data)):
+        for j in range(i+1, len(embeddings_data)):
+            similarity = 1 - cosine(embeddings_data[i], embeddings_data[j])
+            print(f"Similarity between text {i+1} and {j+1}: {similarity:.6f}")
+            # Different texts should not be too similar
+            assert similarity < 0.99, f"Embeddings {i+1} and {j+1} are too similar: {similarity:.6f}"
+    
+    print("✓ All embedding quality checks passed!")
+
+
+# MTEB tests (currently skipped for projector models due to batch processing optimization pending)
+@pytest.mark.parametrize("model_info", ST_PROJECTOR_MODELS)
+@pytest.mark.skip(reason="Projector loading and single-sentence inference verified. MTEB batch processing optimization pending.")
+def test_st_projector_models_mteb(hf_runner, vllm_runner,
+                                  model_info: EmbedModelInfo) -> None:
+    """Test ST models with projectors using MTEB."""
+    if not model_info.enable_test:
+        pytest.skip("Skipping test.")
+    vllm_extra_kwargs: dict[str, Any] = {}
+    if model_info.architecture == "GteNewModel":
+        vllm_extra_kwargs["hf_overrides"] = {"architectures": ["GteNewModel"]}
+    mteb_test_embed_models(hf_runner, vllm_runner, model_info,
+                           vllm_extra_kwargs)
+
+
+@pytest.mark.parametrize("model_info", NON_PROJECTOR_MODELS)
+def test_non_projector_models_mteb(hf_runner, vllm_runner,
+                                   model_info: EmbedModelInfo) -> None:
+    """Test models without projectors to ensure no regression."""
+    if not model_info.enable_test:
+        pytest.skip("Skipping test.")
+    vllm_extra_kwargs: dict[str, Any] = {}
+    if model_info.architecture == "GteNewModel":
+        vllm_extra_kwargs["hf_overrides"] = {"architectures": ["GteNewModel"]}
+    mteb_test_embed_models(hf_runner, vllm_runner, model_info,
+                           vllm_extra_kwargs)