first commit

vllm_br/model_executor/models/internlm2.py

@@ -0,0 +1,140 @@
+################################################################################
+# Copyright(c)2020-2025 Shanghai Biren Technology Co., Ltd. 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.
+#
+################################################################################
+
+# SPDX-License-Identifier: Apache-2.0
+
+from functools import partial
+from typing import Optional, Union
+
+import torch
+
+from vllm.distributed import (get_pp_group, split_tensor_along_last_dim,
+                              tensor_model_parallel_all_gather)
+from vllm.model_executor.layers.activation import SiluAndMul
+from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
+                                               RowParallelLinear)
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.models.internlm2 import (InternLM2Attention,
+                                                  InternLM2MLP, InternLM2Model)
+from vllm.sequence import IntermediateTensors
+
+
+def internlm2_attention_split_qkv(self, qkv: torch.Tensor):
+    seq_len = qkv.shape[1]
+    if self.tp_size > 1:
+        qkv_map = [self.q_size, self.kv_size, self.kv_size] * self.tp_size
+        qkv = tensor_model_parallel_all_gather(qkv)
+        qkv = torch.split(qkv, qkv_map, dim=-1)
+        qkv = qkv[::3] + qkv[1::3] + qkv[2::3]
+        qkv = torch.cat(qkv, dim=-1)
+
+    qkv = qkv.view(seq_len, self.total_num_kv_heads, self.key_value_groups + 2,
+                   self.head_dim)
+    q, k, v = torch.split(qkv, [self.key_value_groups, 1, 1], dim=-2)
+    q = q.reshape(seq_len, self.q_size * self.tp_size).unsqueeze(0)
+    k = k.reshape(seq_len, self.kv_size * self.tp_size).unsqueeze(0)
+    v = v.reshape(seq_len, self.kv_size * self.tp_size).unsqueeze(0)
+
+    if self.tp_size > 1:
+        splitter = partial(split_tensor_along_last_dim,
+                           num_partitions=self.tp_size)
+        q = splitter(q)[self.tp_rank]
+        k = splitter(k)[self.tp_rank]
+        v = splitter(v)[self.tp_rank]
+    return q, k, v
+
+
+def internlm2_attention_forward(
+    self,
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+) -> torch.Tensor:
+    qkv, _ = self.wqkv(hidden_states)
+    q, k, v = self.split_qkv(qkv)
+    q, k = self.rotary_emb(positions, q, k)
+
+    attn_output = self.attn(q, k, v)
+    output, _ = self.wo(attn_output)
+    return output
+
+
+def internlm2_model_forward(
+    self,
+    input_ids: torch.Tensor,
+    positions: torch.Tensor,
+    intermediate_tensors: Optional[IntermediateTensors] = None,
+    inputs_embeds: Optional[torch.Tensor] = None,
+) -> Union[torch.Tensor, IntermediateTensors]:
+    if get_pp_group().is_first_rank:
+        if inputs_embeds is not None:
+            hidden_states = inputs_embeds
+        else:
+            hidden_states = self.get_input_embeddings(input_ids)
+        residual = None
+    else:
+        assert intermediate_tensors is not None
+        hidden_states = intermediate_tensors["hidden_states"]
+        residual = intermediate_tensors["residual"]
+
+    hidden_states = hidden_states.unsqueeze(0)
+
+    for layer in self.layers[self.start_layer:self.end_layer]:
+        hidden_states, residual = layer(positions, hidden_states, residual)
+    if not get_pp_group().is_last_rank:
+        return IntermediateTensors({
+            "hidden_states":
+            hidden_states.squeeze(0) if hidden_states is not None else None,
+            "residual":
+            residual.squeeze(0) if residual is not None else None
+        })
+    hidden_states, _ = self.norm(hidden_states, residual)
+    return hidden_states.squeeze(0)
+
+
+def internlm2_mlp_init(
+    self,
+    hidden_size: int,
+    intermediate_size: int,
+    hidden_act: str,
+    quant_config: Optional[QuantizationConfig] = None,
+    prefix: str = "",
+) -> None:
+    super(InternLM2MLP, self).__init__()
+    self.gate_up_proj = MergedColumnParallelLinear(
+        hidden_size,
+        [intermediate_size] * 2,
+        bias=False,
+        quant_config=quant_config,
+        prefix=f"{prefix}.gate_up_proj",
+    )
+    self.gate_up_proj.no_need_cross = True
+    self.w2 = RowParallelLinear(
+        intermediate_size,
+        hidden_size,
+        bias=False,
+        quant_config=quant_config,
+        prefix=f"{prefix}.w2",
+    )
+    if hidden_act != "silu":
+        raise ValueError(f"Unsupported activation: {hidden_act}. "
+                         "Only silu is supported for now.")
+    self.act_fn = SiluAndMul()
+
+
+InternLM2Attention.split_qkv = internlm2_attention_split_qkv
+InternLM2Attention.forward = internlm2_attention_forward
+InternLM2Model.forward = internlm2_model_forward
+InternLM2MLP.__init__ = internlm2_mlp_init