fix wint8 precision and refine code

Author: yuanlehome

paddlenlp/experimental/transformers/deepseek_v2/modeling.py

@@ -126,28 +126,28 @@ def forward(
         query: paddle.Tensor,
         key: paddle.Tensor,
     ) -> Tuple[paddle.Tensor, paddle.Tensor]:
-        q = query[..., : self.rotary_dim]
-        k = key[..., : self.rotary_dim]
+        query_rot = query[..., : self.rotary_dim]
+        key_rot = key[..., : self.rotary_dim]
         if self.rotary_dim < self.head_size:
             query_pass = query[..., self.rotary_dim :]
             key_pass = key[..., self.rotary_dim :]
         cos_sin = self.cos_sin_cache[position_ids].unsqueeze(1)
         cos, sin = cos_sin.chunk(2, axis=-1)
 
-        s, h, d = q.shape
-        q = q.reshape([s, h, d // 2, 2]).transpose([0, 1, 3, 2]).reshape([s, h, d])
+        s, h, d = query_rot.shape
+        query_rot = query_rot.reshape([s, h, d // 2, 2]).transpose([0, 1, 3, 2]).reshape([s, h, d])
 
-        s, h, d = k.shape
-        k = k.reshape([s, h, d // 2, 2]).transpose([0, 1, 3, 2]).reshape([s, h, d])
+        s, h, d = key_rot.shape
+        key_rot = key_rot.reshape([s, h, d // 2, 2]).transpose([0, 1, 3, 2]).reshape([s, h, d])
 
         def rotate_half(x):
             """Rotates half the hidden axiss of the input."""
             x1 = x[..., : x.shape[-1] // 2]
             x2 = x[..., x.shape[-1] // 2 :]
             return paddle.concat([-x2, x1], axis=-1)  # shape is the same as x
 
-        query_rot = (q * cos) + (rotate_half(q) * sin)
-        key_rot = (k * cos) + (rotate_half(k) * sin)
+        query_rot = (query_rot * cos) + (rotate_half(query_rot) * sin)
+        key_rot = (key_rot * cos) + (rotate_half(key_rot) * sin)
 
         if self.rotary_dim < self.head_size:
             query = paddle.concat((query_rot, query_pass), axis=-1)
@@ -564,6 +564,7 @@ def set_state_dict(self, state_dict):
                         q_b_proj_weight, algo=self.quant_algo
                     )
                     self.transformer_block.q_b_proj_weights[idx].set_value(q_b_proj_quanted_weight)
+                    self.transformer_block.q_a_layernorm_weights[idx].set_value(q_a_layernorm_weight)
                     self.transformer_block.q_b_proj_weights_scale[idx].set_value(q_b_proj_weight_scale)
                 else:
                     self.transformer_block.q_a_proj_weights[idx].set_value(q_a_proj_weight)
@@ -602,6 +603,7 @@ def set_state_dict(self, state_dict):
                     kv_b_proj_weight, algo=self.quant_algo
                 )
                 self.transformer_block.kv_b_proj_weights[idx].set_value(kv_b_proj_quanted_weight)
+                self.transformer_block.kv_a_layernorm_weights[idx].set_value(kv_a_layernorm_weight)
                 self.transformer_block.kv_b_proj_weights_scale[idx].set_value(kv_b_proj_weight_scale)
             else:
                 self.transformer_block.kv_a_proj_with_mqa_weights[idx].set_value(kv_a_proj_with_mqa_weight)

paddlenlp/experimental/transformers/fused_transformer_layers.py

@@ -415,6 +415,8 @@ def __init__(self, config: FusedMultiTransformerConfig):
             mscale = self.config.mla_config.mscale
             self.softmax_scale = self.softmax_scale * mscale * mscale
 
+        self.position_ids: list[int] = []
+
         self.weight_dtype = self._dtype
         self.create_params_type = self.get_weight_create_dype()
 
@@ -949,7 +951,7 @@ def compute_layernorm_before_qkv(self, src, i):
 
         return ln_out
 
-    def compute_qkv_linear(self, ln_out, i, position_ids=None):
+    def compute_qkv_linear(self, ln_out, i):
         if self.config.mla_config.use_mla():
             if self.config.mla_config.q_lora_rank is not None:
                 query = paddle.matmul(ln_out, self.q_a_proj_weights[i])
@@ -989,7 +991,7 @@ def compute_qkv_linear(self, ln_out, i, position_ids=None):
                 key_value, [self.config.mla_config.qk_nope_head_dim, self.config.mla_config.v_head_dim], axis=-1
             )
 
-            query_pe, key_pe = self.config.rotary_emb(position_ids, query_pe, key_pe)
+            query_pe, key_pe = self.config.rotary_emb(self.position_ids, query_pe, key_pe)
 
             query[..., self.config.mla_config.qk_nope_head_dim :] = query_pe
             key = paddle.empty_like(query)
@@ -1017,9 +1019,9 @@ def compute_qkv_linear(self, ln_out, i, position_ids=None):
 
         return qkv_out
 
-    def compute_qkv(self, src, residual_input, i, position_ids=None):
+    def compute_qkv(self, src, residual_input, i):
         ln_out = self.compute_layernorm_before_qkv(src, i)
-        qkv_out = self.compute_qkv_linear(ln_out, i, position_ids)
+        qkv_out = self.compute_qkv_linear(ln_out, i)
         return qkv_out, residual_input
 
     def compute_max_len(self, seq_lens_encoder, seq_lens_decoder, cum_offsets):
@@ -1298,7 +1300,20 @@ def compute_shared_expert(self, tmp_out, i):
         return ffn2_out
 
     def pre_process(self, **kwargs):
-        pass
+        seq_lens_this_time = kwargs.get("seq_lens_this_time", None)
+        bsz = seq_lens_this_time.shape[0]
+        position_ids = []
+        for i in range(bsz):
+            cur_seq_len = kwargs.get("seq_lens_encoder", None)[i]
+            if cur_seq_len > 0:
+                for j in range(cur_seq_len):
+                    position_ids.append(j)
+            else:
+                ids = kwargs.get("seq_lens_decoder", None)[i].item()
+                if ids > 0:
+                    position_ids.append(ids)
+
+        self.position_ids = position_ids
 
     def post_process(self, **kwargs):
         time_step = kwargs.get("time_step", None)
@@ -1405,23 +1420,10 @@ def forward(
                 kwargs.get("block_size", 64),
                 self.config.speculate_config.speculate_max_draft_token_num,
             )
-        seq_lens_this_time = kwargs.get("seq_lens_this_time", None)
-        bsz = seq_lens_this_time.shape[0]
-        position_ids = []
-        for i in range(bsz):
-            cur_seq_len = kwargs.get("seq_lens_encoder", None)[i]
-            if cur_seq_len > 0:
-                for j in range(cur_seq_len):
-                    position_ids.append(j)
-            else:
-                ids = kwargs.get("seq_lens_decoder", None)[i].item()
 
-                if ids > 0:
-                    position_ids.append(ids)
-        # print("position_ids;", position_ids)
         residual_input = src
         for i in range(self.num_layers):
-            qkv_out, residual_input = self.compute_qkv(src, residual_input, i, position_ids)
+            qkv_out, residual_input = self.compute_qkv(src, residual_input, i)
             out_linear_out = self.compute_attn(
                 time_step,
                 qkv_out,
@@ -1856,8 +1858,7 @@ def compute_qkv_linear(self, ln_out, i):
                 key_value, [self.config.mla_config.qk_nope_head_dim, self.config.mla_config.v_head_dim], axis=-1
             )
 
-            position_ids = paddle.arange(ln_out.shape[0]).reshape((1, -1))
-            query_pe, key_pe = self.config.rotary_emb(position_ids, query_pe, key_pe)
+            query_pe, key_pe = self.config.rotary_emb(self.position_ids, query_pe, key_pe)
 
             query[..., self.config.mla_config.qk_nope_head_dim :] = query_pe
             key = paddle.empty_like(query)