[Feature] support qwen3-embedding model load (#4202)

* support qwen3-embedding

* fix ci bug

* fix

* fix ci bug

* fix ci bug

* fix

Author: lizexu123

fastdeploy/model_executor/layers/embeddings.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 """
 
+from dataclasses import dataclass
 from typing import Dict
 
 import numpy as np
@@ -22,9 +23,73 @@
 from paddle.distributed import fleet
 
 from fastdeploy.config import FDConfig
-from fastdeploy.model_executor.utils import set_weight_attrs
+from fastdeploy.model_executor.utils import set_weight_attrs, slice_fn
 
-from .utils import get_tensor
+from .utils import (
+    DEFAULT_VOCAB_PADDING_SIZE,
+    get_tensor,
+    pad_vocab_size,
+    vocab_range_from_global_vocab_size,
+)
+
+
+@dataclass
+class VocabParallelEmbeddingShardIndices:
+    """Indices for a shard of a vocab parallel embedding."""
+
+    padded_org_vocab_start_index: int
+    padded_org_vocab_end_index: int
+    padded_added_vocab_start_index: int
+    padded_added_vocab_end_index: int
+
+    org_vocab_start_index: int
+    org_vocab_end_index: int
+    added_vocab_start_index: int
+    added_vocab_end_index: int
+
+    @property
+    def num_org_elements(self) -> int:
+        return self.org_vocab_end_index - self.org_vocab_start_index
+
+    @property
+    def num_added_elements(self) -> int:
+        return self.added_vocab_end_index - self.added_vocab_start_index
+
+    @property
+    def num_org_elements_padded(self) -> int:
+        return self.padded_org_vocab_end_index - self.padded_org_vocab_start_index
+
+    @property
+    def num_added_elements_padded(self) -> int:
+        return self.padded_added_vocab_end_index - self.padded_added_vocab_start_index
+
+    @property
+    def num_org_vocab_padding(self) -> int:
+        return self.num_org_elements_padded - self.num_org_elements
+
+    @property
+    def num_added_vocab_padding(self) -> int:
+        return self.num_added_elements_padded - self.num_added_elements
+
+    @property
+    def num_elements_padded(self) -> int:
+        return self.num_org_elements_padded + self.num_added_elements_padded
+
+    def __post_init__(self):
+        # sanity checks
+        assert self.padded_org_vocab_start_index <= self.padded_org_vocab_end_index
+        assert self.padded_added_vocab_start_index <= self.padded_added_vocab_end_index
+
+        assert self.org_vocab_start_index <= self.org_vocab_end_index
+        assert self.added_vocab_start_index <= self.added_vocab_end_index
+
+        assert self.org_vocab_start_index <= self.padded_org_vocab_start_index
+        assert self.added_vocab_start_index <= self.padded_added_vocab_start_index
+        assert self.org_vocab_end_index <= self.padded_org_vocab_end_index
+        assert self.added_vocab_end_index <= self.padded_added_vocab_end_index
+
+        assert self.num_org_elements <= self.num_org_elements_padded
+        assert self.num_added_elements <= self.num_added_elements_padded
 
 
 class VocabParallelEmbedding(nn.Layer):
@@ -39,6 +104,7 @@ def __init__(
         embedding_dim: int = 768,
         params_dtype: str = "bfloat16",
         prefix="",
+        padding_size: int = DEFAULT_VOCAB_PADDING_SIZE,
     ) -> None:
         """
         Initialize the VocabParallelEmbedding layer for the model.
@@ -65,18 +131,40 @@ def __init__(
         self.max_position_embeddings: int = fd_config.model_config.max_position_embeddings
         self.tie_word_embeddings: bool = fd_config.model_config.tie_word_embeddings
         self.params_dtype: str = params_dtype
+        self.padding_size = padding_size
+
+        self.org_vocab_size = num_embeddings
+        self.num_embeddings = num_embeddings
+        num_added_embeddings = num_embeddings - self.org_vocab_size
+
+        self.org_vocab_size_padded = pad_vocab_size(self.org_vocab_size, self.padding_size)
+        self.num_embeddings_padded = pad_vocab_size(
+            self.org_vocab_size_padded + num_added_embeddings, self.padding_size
+        )
+        assert self.org_vocab_size_padded <= self.num_embeddings_padded
+        self.shard_indices = self._get_indices(
+            self.num_embeddings_padded,
+            self.org_vocab_size_padded,
+            self.num_embeddings,
+            self.org_vocab_size,
+            self.tensor_parallel_rank,
+            self.world_size,
+        )
+
+        if num_embeddings % self.world_size != 0:
+            self.num_embeddings_padded = pad_vocab_size(num_embeddings, self.padding_size)
 
         if not self.column_cut:
             self.embeddings = fleet.meta_parallel.VocabParallelEmbedding(
-                num_embeddings,
+                self.num_embeddings_padded,
                 embedding_dim,
                 mp_group=self.tp_group,
                 weight_attr=paddle.ParamAttr(
                     initializer=nn.initializer.Normal(mean=0.0, std=self.initializer_range),
                 ),
             )
             if self.world_size > 1:
-                set_weight_attrs(self.embeddings.weight, {"output_dim": False})
+                set_weight_attrs(self.embeddings.weight, {"output_dim": False, "weight_loader": self.weight_loader})
         else:
             # column cut embedding
             self.embeddings = nn.Embedding(
@@ -106,6 +194,88 @@ def load_state_dict(self, state_dict: Dict[str, paddle.Tensor | np.ndarray]):
 
         self.embeddings.weight.set_value(weight_tensor)
 
+    @classmethod
+    def _get_indices(
+        cls,
+        vocab_size_paded: int,
+        org_vocab_size_padded: int,
+        vocab_size: int,
+        org_vocab_size: int,
+        tp_rank: int,
+        tp_size: int,
+    ) -> VocabParallelEmbeddingShardIndices:
+        """Get start and end indices for vocab parallel embedding, following the
+        layout outlined in the class docstring, based on the given tp_rank and
+        tp_size."""
+
+        num_added_embeddings_padded = vocab_size_paded - org_vocab_size_padded
+        padded_org_vocab_start_index, padded_org_vocab_end_index = vocab_range_from_global_vocab_size(
+            org_vocab_size_padded, tp_rank, tp_size
+        )
+
+        padded_added_vocab_start_index, padded_added_vocab_end_index = vocab_range_from_global_vocab_size(
+            num_added_embeddings_padded, tp_rank, tp_size, offset=org_vocab_size
+        )
+        # remove padding
+        org_vocab_start_index = min(padded_org_vocab_start_index, org_vocab_size)
+        org_vocab_end_index = min(padded_org_vocab_end_index, org_vocab_size)
+        added_vocab_start_index = min(padded_added_vocab_start_index, vocab_size)
+        added_vocab_end_index = min(padded_added_vocab_end_index, vocab_size)
+        return VocabParallelEmbeddingShardIndices(
+            padded_org_vocab_start_index,
+            padded_org_vocab_end_index,
+            padded_added_vocab_start_index,
+            padded_added_vocab_end_index,
+            org_vocab_start_index,
+            org_vocab_end_index,
+            added_vocab_start_index,
+            added_vocab_end_index,
+        )
+
+    def weight_loader(self, param, loaded_weight, shard_id=None):
+        output_dim = getattr(param, "output_dim", None)
+        packed_dim = getattr(param, "packed_dim", None)
+
+        loaded_weight = get_tensor(loaded_weight)
+        if param.dtype != loaded_weight.dtype:
+            if loaded_weight.dtype == paddle.int8 and param.dtype == paddle.float8_e4m3fn:
+                loaded_weight = loaded_weight.cast(param.dtype)
+            else:
+                loaded_weight = loaded_weight.cast(param.dtype)
+
+        if output_dim is None:
+            assert (
+                param.shape == loaded_weight.shape
+            ), f"Shape mismatch: param {param.shape} vs loaded_weight {loaded_weight.shape}"
+            param.set_value(loaded_weight)
+            return
+
+        start_idx = self.shard_indices.org_vocab_start_index
+        end_idx = self.shard_indices.org_vocab_end_index
+        shard_size = self.shard_indices.org_vocab_end_index - start_idx
+
+        # If param packed on the same dim we are sharding on, then
+        # need to adjust offsets of loaded weight by pack_factor.
+        if packed_dim is not None and packed_dim == output_dim:
+            packed_factor = getattr(param, "packed_factor", getattr(param, "pack_factor", 1))
+            assert loaded_weight.shape[output_dim] == (self.org_vocab_size // packed_factor)
+            start_idx = start_idx // packed_factor
+            shard_size = shard_size // packed_factor
+        else:
+            assert loaded_weight.shape[output_dim] == self.org_vocab_size, (
+                f"Loaded weight dim {output_dim} size {loaded_weight.shape[output_dim]} "
+                f"!= org_vocab_size {self.org_vocab_size}"
+            )
+
+        shard_weight = slice_fn(loaded_weight, output_dim, start_idx, end_idx)
+
+        if output_dim == 0:
+            param[: shard_weight.shape[0]].copy_(shard_weight, False)
+            param[shard_weight.shape[0] :].fill_(0)
+        else:
+            param[:, : shard_weight.shape[1]].copy_(shard_weight, False)
+            param[:, shard_weight.shape[1] :].fill_(0)
+
     def forward(self, ids_remove_padding=None) -> paddle.Tensor:
         """
         Defines the forward computation of the layer.

fastdeploy/model_executor/layers/lm_head.py

@@ -22,6 +22,10 @@
 from paddle.distributed import fleet
 
 from fastdeploy.config import FDConfig
+from fastdeploy.model_executor.layers.utils import (
+    DEFAULT_VOCAB_PADDING_SIZE,
+    pad_vocab_size,
+)
 from fastdeploy.model_executor.utils import (
     default_weight_loader,
     set_weight_attrs,
@@ -44,6 +48,7 @@ def __init__(
         prefix: str = "",
         with_bias: bool = False,
         dtype: str = None,
+        padding_size: int = DEFAULT_VOCAB_PADDING_SIZE,
     ) -> None:
         """
         Parallelized LMhead.
@@ -68,6 +73,10 @@ def __init__(
         self.column_cut = True
         self.nranks = fd_config.parallel_config.tensor_parallel_size
         self.fd_config = fd_config
+        self.padding_size = padding_size
+
+        if num_embeddings % self.nranks != 0:
+            num_embeddings = pad_vocab_size(num_embeddings, self.padding_size)
 
         ColumnParallelLinear = fleet.meta_parallel.ColumnParallelLinear
         RowParallelLinear = fleet.meta_parallel.RowParallelLinear

fastdeploy/model_executor/layers/pool/__init__.py

@@ -0,0 +1,15 @@
+"""
+# Copyright (c) 2025  PaddlePaddle Authors. All Rights Reserved.
+#
+# 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.
+"""

fastdeploy/model_executor/layers/utils.py

@@ -45,6 +45,14 @@
     c8_state_dict = paddle.load(cache_params, return_numpy=True)
 
 
+DEFAULT_VOCAB_PADDING_SIZE = 64
+
+
+def pad_vocab_size(vocab_size: int, pad_to: int = DEFAULT_VOCAB_PADDING_SIZE) -> int:
+    """Pad the vocab size to the given value."""
+    return ((vocab_size + pad_to - 1) // pad_to) * pad_to
+
+
 def per_block_cast_to_fp8(x: Tensor, block_size: list = [128, 128]) -> Tuple[Tensor, Tensor]:
     """
     Only used in deep_gemm block wise quant weight.
@@ -372,3 +380,14 @@ def create_empty_tensor(shape: Tuple[int, ...], dtype: Union[paddle.dtype, str])
         paddle.Tensor: An empty tensor with the specified shape and data type.
     """
     return paddle.empty(list(shape), dtype=dtype)
+
+
+def vocab_range_from_per_partition_vocab_size(per_partition_vocab_size: int, rank: int, offset: int = 0):
+    index_f = rank * per_partition_vocab_size
+    index_l = index_f + per_partition_vocab_size
+    return index_f + offset, index_l + offset
+
+
+def vocab_range_from_global_vocab_size(global_vocab_size: int, rank: int, world_size: int, offset: int = 0):
+    per_partition_vocab_size = divide(global_vocab_size, world_size)
+    return vocab_range_from_per_partition_vocab_size(per_partition_vocab_size, rank, offset=offset)