fix precision(ongoing)

Author: shihaobai

lightllm/models/vit/layer_infer/transformer_layer_infer.py

@@ -31,18 +31,22 @@ def __init__(self, layer_num, tp_rank, world_size, network_config, mode=[]):
 
     def norm(self, input, weight):
         input_dtype = input.dtype
+        input_shape = input.shape
+        input = input.view(-1, self.tp_padding_head_num * self.head_dim_)
         input = input.to(torch.float32)
         variance = input.pow(2).mean(-1, keepdim=True)
         input = input * torch.rsqrt(variance + self.eps_)
-        return weight * input.to(input_dtype)
+        out = weight * input.to(input_dtype)
+        out = out.reshape(input_shape)
+        return out
 
     def tp_norm(self, input, weight):
         input_shape = input.shape
         input = input.view(-1, self.tp_padding_head_num * self.head_dim_)
         input_dtype = input.dtype
         input = input.to(torch.float32)
         tp_variance = input.pow(2).sum(-1, keepdim=True)
-        dist.all_reduce(tp_variance, op=dist.ReduceOp.SUM, async_op=False)
+        # dist.all_reduce(tp_variance, op=dist.ReduceOp.SUM, async_op=False)
         variance = tp_variance / self.embed_dim_
         input = input * torch.rsqrt(variance + self.eps_)
         out = weight * input.to(input_dtype)
@@ -75,8 +79,8 @@ def _ffn_norm(self, input, layer_weight: ViTTransformerLayerWeight) -> torch.Ten
             )
 
     def _qk_norm(self, q, k, layer_weight: ViTTransformerLayerWeight) -> torch.Tensor:
-        q_norm = self.tp_norm(q, layer_weight.q_norm_weight_.weight)
-        k_norm = self.tp_norm(k, layer_weight.k_norm_weight_.weight)
+        q_norm = self.norm(q, layer_weight.q_norm_weight_.weight)
+        k_norm = self.norm(k, layer_weight.k_norm_weight_.weight)
         return q_norm, k_norm
 
     def _get_qkv(self, input, layer_weight: ViTTransformerLayerWeight) -> torch.Tensor:
@@ -85,6 +89,9 @@ def _get_qkv(self, input, layer_weight: ViTTransformerLayerWeight) -> torch.Tens
         qkv = layer_weight.qkv_proj.mm(input.view(-1, self.embed_dim_), use_custom_tensor_mananger=False)
         qkv = qkv.view(batch_size, seq_len, 3, -1, self.head_dim_)
         q, k, v = qkv.unbind(2)
+        q = q.contiguous()
+        k = k.contiguous()
+        v = v.contiguous()
         return q, k, v
 
     def _context_attention_kernel(self, q, k, v) -> torch.Tensor:

lightllm/models/vit/layer_weights/transformer_layer_weight.py

@@ -163,5 +163,6 @@ def load_hf_weights(self, weights):
             ls2 = weights[f"vision_model.encoder.layers.{self.layer_num_}.ls2"]
             self.ls2 = self._cuda(ls2)
             self.use_ls = True
+            print(self.ls1)
 
         return super().load_hf_weights(weights)

lightllm/models/vit/model.py

@@ -163,6 +163,7 @@ def encode(self, image_uuids: List):
 
         imgs = torch.cat(img_tensors, dim=0)
         pixel_values = imgs.cuda().to(dtype=self.data_type)
+        print(pixel_values.shape, pixel_values.dtype)
         all_img_embeds = self.forward(pixel_values)
         return all_img_embeds, uuids, valid_ids
 

lightllm/server/visualserver/model_infer/model_rpc.py

@@ -63,6 +63,7 @@ def exposed_init_model(self, kvargs):
                     "quant_cfg": kvargs["quant_cfg"],
                 }
                 self.model = VisionTransformer(kvargs)
+                # self.model = InternVLVisionModel()
             else:
                 raise Exception(f"can not support {self.model_type} now")