optimizer_dual_pp_post_node_memory (#10806)

Author: phlrain

paddlenlp/transformers/deepseek_v2/modeling_pp.py

@@ -153,13 +153,10 @@ def forward(self, inputs):
 
         with paddle.no_grad():
             if self.shared_experts is not None:
-                x_fp8, x_scale, shared_expert_output = fp8_mlp_fwd(
-                    hidden_states, self.shared_experts.w1, self.shared_experts.w2
-                )
+                shared_expert_output = fp8_mlp_fwd(hidden_states, self.shared_experts.w1, self.shared_experts.w2)
                 final_hidden_states = final_hidden_states + shared_expert_output
 
-        self.x_fp8 = x_fp8
-        self.x_scale = x_scale
+        self.x = hidden_states
         self.l_aux = l_aux
         hidden_states = residual + final_hidden_states
 
@@ -174,10 +171,9 @@ def backward(self, output_grad):
 
         assert not self.send_mtp_embed, "not support have mtp have yet"
 
-        dx = fp8_mlp_bwd(do3, self.x_fp8, self.x_scale, self.shared_experts.w1, self.shared_experts.w2)
+        dx = fp8_mlp_bwd(do3, self.x, self.shared_experts.w1, self.shared_experts.w2)
 
-        self.x_fp8 = None
-        self.x_scale = None
+        self.x = None
 
         residual_grad = do3
 

paddlenlp/transformers/fp8_utils.py

@@ -284,28 +284,35 @@ def fp8_mlp_fwd(x, w1, w2):
     if len(x_orig_shape) > 2:
         o3 = o3.reshape([x_orig_shape[0], -1, o3.shape[-1]])
 
-    return x_fp8, x_scale, o3
+    return o3
 
 
-def fp8_mlp_bwd(do3, x_fp8, x_scale, w1, w2):
+def fp8_mlp_bwd(do3, x, w1, w2):
     do3_orig_shape = do3.shape
     do3 = do3.reshape([-1, do3_orig_shape[-1]])
 
-    x_orig_shape = x_fp8.shape
+    x_orig_shape = x.shape
+    x = x.reshape([-1, x_orig_shape[-1]])
+
+    if x.shape[0] % 128 == 0:
+        x_fp8, x_scale, x_t_fp8, x_t_scale = paddle.incubate.nn.functional.fp8_quant_blockwise(
+            x, output_scale_transpose=True, quant_method="1x128", input_transpose=True
+        )
+    else:
+        x_fp8, x_scale = paddle.incubate.nn.functional.fp8_quant_blockwise(
+            x, output_scale_transpose=True, quant_method="1x128", input_transpose=False
+        )
+        x = padding(x, 0)
+        _, _, x_t_fp8, x_t_scale = paddle.incubate.nn.functional.fp8_quant_blockwise(
+            x, output_scale_transpose=True, quant_method="1x128", input_transpose=True
+        )
 
     _, _, w1_fp8, w1_sacle = paddle.incubate.nn.functional.fp8_quant_blockwise(
         w1, output_scale_transpose=False, quant_method="128x128", input_transpose=True
     )
     o1 = paddle.empty([x_fp8.shape[0], w1_fp8.shape[0]], dtype=do3.dtype)
     deep_gemm.gemm_fp8_fp8_bf16_nt((x_fp8, x_scale.T), (w1_fp8, w1_sacle), o1, num_sms=112)
 
-    x_dequant_fp16 = paddle.incubate.nn.functional.fused_act_dequant(x_fp8, x_scale.T.contiguous())
-    x_dequant_fp16 = padding(x_dequant_fp16, 0)
-
-    _, _, x_t_fp8, x_t_scale = paddle.incubate.nn.functional.fp8_quant_blockwise(
-        x_dequant_fp16, output_scale_transpose=True, quant_method="1x128", input_transpose=True
-    )
-
     # ===== [recompute] o2 = swiglu(o1) =====
     o2 = swiglu(o1)
 
@@ -577,7 +584,10 @@ def split_group_gemm(x_fp8, x_scale, w_fp8, w_scale, tokens_per_expert, gemm_out
         x_scale_tma_align = x_scale[start_idx:end_idx].T.contiguous().T
 
         deep_gemm.gemm_fp8_fp8_bf16_nt(
-            (x_fp8[start_idx:end_idx], x_scale_tma_align), (w_fp8[i], w_scale[i]), gemm_out[start_idx:end_idx], num_sms=112
+            (x_fp8[start_idx:end_idx], x_scale_tma_align),
+            (w_fp8[i], w_scale[i]),
+            gemm_out[start_idx:end_idx],
+            num_sms=112,
         )
 
         start_idx = end_idx
@@ -763,7 +773,7 @@ def bwd_dowm_input(self, expert_w2, unzipped_grad, o1, inplace_swiglu_prob=False
                     (bw_w2_quant, bw_w2_scale),
                     do2_s,
                     m_indices=self.m_indices,
-                    num_sms=112
+                    num_sms=112,
                 )
 
         with paddle.amp.auto_cast(False):