refine code

Author: yuanlehome

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)