Update vllm_plugin/cloverlm_vllm.py

vllm_plugin/cloverlm_vllm.py

@@ -12,6 +12,7 @@ from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
+    ReplicatedLinear,
     RowParallelLinear,
 )
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -22,6 +23,10 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
 )
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper
+from vllm.distributed import (
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+)
 
 
 def _build_rope_cos_sin(
@@ -62,45 +67,64 @@ class CloverLMAttention(nn.Module):
         prefix: str = "",
     ):
         super().__init__()
-        self.num_heads = num_heads
-        self.num_kv_heads = num_kv_heads
+        tp_size = get_tensor_model_parallel_world_size()
+        tp_rank = get_tensor_model_parallel_rank()
+
+        self.num_heads = num_heads // tp_size
         self.head_dim = head_dim
-        self.q_size = num_heads * head_dim
-        self.kv_size = num_kv_heads * head_dim
+        self.q_size = self.num_heads * head_dim
+
+        total_q_size = num_heads * head_dim
+        total_kv_size = num_kv_heads * head_dim
+
+        if num_kv_heads % tp_size == 0:
+            self.num_kv_heads = num_kv_heads // tp_size
+            kv_linear_cls = ColumnParallelLinear
+        else:
+            self.num_kv_heads = num_kv_heads
+            kv_linear_cls = ReplicatedLinear
+
+        self.kv_size = self.num_kv_heads * head_dim
 
         self.lq = ColumnParallelLinear(
-            d, self.q_size, bias=False,
+            d, total_q_size, bias=False,
             quant_config=quant_config,
             prefix=f"{prefix}.lq",
         )
-        self.lk = ColumnParallelLinear(
-            d, self.kv_size, bias=False,
+        self.lk = kv_linear_cls(
+            d, total_kv_size, bias=False,
             quant_config=quant_config,
             prefix=f"{prefix}.lk",
         )
-        self.lv = ColumnParallelLinear(
-            d, self.kv_size, bias=False,
+        self.lv = kv_linear_cls(
+            d, total_kv_size, bias=False,
             quant_config=quant_config,
             prefix=f"{prefix}.lv",
         )
         self.lo = RowParallelLinear(
-            self.q_size, d, bias=False,
+            total_q_size, d, bias=False,
             quant_config=quant_config,
             prefix=f"{prefix}.lo",
         )
 
-        # Per-head learnable scale: stored as (1, heads, 1, 1) in checkpoint,
-        # reshaped to (heads,) for efficient multiply after sphere norm.
         self.scale = nn.Parameter(
-            torch.empty(1, num_heads, 1, 1),
+            torch.empty(1, self.num_heads, 1, 1),
             requires_grad=False,
         )
+        heads_per_tp = self.num_heads
+
+        def _scale_weight_loader(param, loaded_weight):
+            start = tp_rank * heads_per_tp
+            end = start + heads_per_tp
+            param.data.copy_(loaded_weight[:, start:end, :, :])
+
+        self.scale.weight_loader = _scale_weight_loader
 
         self.attn = Attention(
-            num_heads=num_heads,
+            num_heads=self.num_heads,
             head_size=head_dim,
             scale=1.0,
-            num_kv_heads=num_kv_heads,
+            num_kv_heads=self.num_kv_heads,
             cache_config=cache_config,
             quant_config=quant_config,
             prefix=f"{prefix}.attn",