fix vit-6b

Author: shihaobai

lightllm/models/vit/layer_infer/transformer_layer_infer.py

@@ -89,16 +89,8 @@ def _ffn_norm(self, input, layer_weight: ViTTransformerLayerWeight) -> torch.Ten
             )
 
     def _qk_norm(self, q, k, layer_weight: ViTTransformerLayerWeight) -> torch.Tensor:
-        if self.tp_world_size_ > 1:
-            q_norm = self.tp_norm(q, layer_weight.q_norm_weight_.weight)
-            k_norm = self.tp_norm(k, layer_weight.k_norm_weight_.weight)
-        else:
-            q_norm = rms_norm(
-                q, weight=layer_weight.q_norm_weight_.weight, eps=self.eps_, use_custom_tensor_mananger=True
-            )
-            k_norm = rms_norm(
-                k, weight=layer_weight.k_norm_weight_.weight, eps=self.eps_, use_custom_tensor_mananger=True
-            )
+        q_norm = self.tp_norm(q, layer_weight.q_norm_weight_.weight)
+        k_norm = self.tp_norm(k, layer_weight.k_norm_weight_.weight)
         return q_norm, k_norm
 
     def _get_qkv(self, input, layer_weight: ViTTransformerLayerWeight) -> torch.Tensor:

lightllm/models/vit/triton_kernel/rms_norm_vit.py

@@ -35,6 +35,7 @@ def rms_norm_kernel(
 
 def rms_norm(hidden_states: Tensor, weight: Tensor, eps: float = 1e-5, use_custom_tensor_mananger: bool = False):
     """Rms norm."""
+    assert len(hidden_states.shape) == 2
     feat_size = weight.shape[0]
     seq_len = hidden_states.numel() // hidden_states.size(-1)
     input_stride = hidden_states.stride(-2)