fix pp hook

Author: eee4017

python/paddle/distributed/fleet/meta_parallel/pipeline_parallel.py

@@ -594,7 +594,9 @@ def forward_backward_pipeline(
             )
 
             self._record_stamp("F", step_id, '"B"', self._forward_color)
-            output_tensor = self._forward_step(input_tensor, micro_dataset)
+            output_tensor = self._forward_step(
+                input_tensor, micro_dataset, step_id=step_id
+            )
             self._record_stamp("F", step_id, '"E"', self._forward_color)
             self._p2p_helper.send_forward(
                 output_tensor,
@@ -626,7 +628,9 @@ def forward_backward_pipeline(
             self._record_stamp(
                 "F", startup_steps + i, '"B"', self._forward_color
             )
-            output_tensor = self._forward_step(input_tensor, micro_dataset)
+            output_tensor = self._forward_step(
+                input_tensor, micro_dataset, step_id=startup_steps + i
+            )
             self._record_stamp(
                 "F", startup_steps + i, '"E"', self._forward_color
             )
@@ -649,7 +653,7 @@ def forward_backward_pipeline(
 
             self._record_stamp("B", i, '"B"', self._backward_color)
             input_tensor_grad = self._backward_step(
-                input_tensor, output_tensor, output_tensor_grad
+                input_tensor, output_tensor, output_tensor_grad, step_id=i
             )
             self._record_stamp("B", i, '"E"', self._backward_color)
 
@@ -684,7 +688,10 @@ def forward_backward_pipeline(
                 "B", steady_steps + i, '"B"', self._backward_color
             )
             input_tensor_grad = self._backward_step(
-                input_tensor, output_tensor, output_tensor_grad
+                input_tensor,
+                output_tensor,
+                output_tensor_grad,
+                step_id=steady_steps + i,
             )
             self._record_stamp(
                 "B", steady_steps + i, '"E"', self._backward_color
@@ -844,7 +851,9 @@ def eval_batch(self, data, compute_loss=False, loss_fn_idx=0):
                 self.is_pipeline_first_stage()
             )
 
-            output_tensor = self._forward_step(input_tensor, micro_dataset)
+            output_tensor = self._forward_step(
+                input_tensor, micro_dataset, step_id=None
+            )
             self._p2p_helper.send_forward(
                 output_tensor,
                 self.is_pipeline_last_stage(),
@@ -862,7 +871,9 @@ def eval_batch(self, data, compute_loss=False, loss_fn_idx=0):
         for i in range(steady_steps):
             last_iter = i == (steady_steps - 1)
 
-            output_tensor = self._forward_step(input_tensor, micro_dataset)
+            output_tensor = self._forward_step(
+                input_tensor, micro_dataset, step_id=None
+            )
             self._p2p_helper.send_forward(
                 output_tensor,
                 self.is_pipeline_last_stage(),
@@ -884,7 +895,9 @@ def eval_batch(self, data, compute_loss=False, loss_fn_idx=0):
 
         return self.train_loss
 
-    def _forward_step(self, input_tensor, micro_dataset, chunk_id=None):
+    def _forward_step(
+        self, input_tensor, micro_dataset, chunk_id=None, step_id=None
+    ):
         if self._enable_timer:
             self.timers("forward_step").start()
         if self.is_pipeline_first_stage():
@@ -893,7 +906,18 @@ def _forward_step(self, input_tensor, micro_dataset, chunk_id=None):
 
         assert chunk_id is None or isinstance(chunk_id, int)
 
+        self.callbacks.on_location(
+            PipelineParallelMicroStepLocations.FORWARD_BEGIN,
+            input_tensor=input_tensor,
+            step_id=step_id,
+        )
         output_tensor = self._layers.forward(input_tensor, chunk_id=chunk_id)
+        self.callbacks.on_location(
+            PipelineParallelMicroStepLocations.FORWARD_END,
+            input_tensor=input_tensor,
+            output_tensor=output_tensor,
+            step_id=step_id,
+        )
 
         if self.is_pipeline_last_stage():
             # train calculate loss for train
@@ -935,10 +959,19 @@ def _forward_step(self, input_tensor, micro_dataset, chunk_id=None):
             return backward_loss_tensor
         return output_tensor
 
-    def _backward_step(self, input_tensor, output_tensor, output_tensor_grad):
+    def _backward_step(
+        self, input_tensor, output_tensor, output_tensor_grad, step_id=None
+    ):
         if self._enable_timer:
             self.timers("backward_step").start()
         with paddle.amp.auto_cast(enable=False):
+            self.callbacks.on_location(
+                PipelineParallelMicroStepLocations.BACKWARD_BEGIN,
+                input_tensor=input_tensor,
+                output_tensor=output_tensor,
+                output_tensor_grad=output_tensor_grad,
+                step_id=step_id,
+            )
             if self.is_pipeline_last_stage():
                 assert output_tensor_grad is None
                 if self.scaler:
@@ -969,6 +1002,14 @@ def _backward_step(self, input_tensor, output_tensor, output_tensor_grad):
                     input_tensor_grad = input_tensor.grad
             if self._enable_timer:
                 self.timers("backward_step").stop()
+            self.callbacks.on_location(
+                PipelineParallelMicroStepLocations.BACKWARD_END,
+                input_tensor=input_tensor,
+                output_tensor=output_tensor,
+                input_tensor_grad=input_tensor_grad,
+                output_tensor_grad=output_tensor_grad,
+                step_id=step_id,
+            )
             return input_tensor_grad
 
     def _check_micro_batch_data_valid(self, micro_batch_data):
@@ -1217,7 +1258,7 @@ def _forward_step_helper(self, micro_dataset, micro_step):
         )
         input_tensor = self.input_tensors[virtual_pp_rank][-1]
         output_tensor = self._forward_step(
-            input_tensor, micro_dataset, virtual_pp_rank
+            input_tensor, micro_dataset, virtual_pp_rank, step_id=micro_step
         )
         self.output_tensors[virtual_pp_rank].append(output_tensor)
 
@@ -1281,7 +1322,7 @@ def _backward_step_helper(self, micro_step):
         output_tensor = self.output_tensors[virtual_pp_rank].pop(0)
         output_tensor_grad = self.output_tensor_grads[virtual_pp_rank].pop(0)
         input_tensor_grad = self._backward_step(
-            input_tensor, output_tensor, output_tensor_grad
+            input_tensor, output_tensor, output_tensor_grad, step_id=micro_step
         )
 
         self._overlap_comm_grads()