supoort_profile_data_load

Author: xuxinyi389

paddleformers/trainer/trainer.py

@@ -1924,7 +1924,15 @@ def _inner_training_loop(
 
             step = -1
 
+            # Data loading timing for global_step
+            _data_load_time_for_global_step = 0.0
+            _data_load_start_time = time.time()
+
             for step, inputs in enumerate(epoch_iterator):
+                # Record data loading time for this iteration
+                _data_load_end_time = time.time()
+                _data_load_time_for_global_step += _data_load_end_time - _data_load_start_time
+
                 if self.args.profile and step % self.args.gradient_accumulation_steps == 0:
                     perf_utils.switch_profile(
                         self.state.global_step,
@@ -1959,6 +1967,8 @@ def _inner_training_loop(
                     if steps_trained_in_current_epoch == 0:
                         self._load_rng_state(resume_from_checkpoint)
                     self.timers and self.timers("read-data").start()
+                    # Reset data loading timer for skipped steps
+                    _data_load_start_time = time.time()
                     continue
                 elif steps_trained_progress_bar is not None:
                     steps_trained_progress_bar.close()
@@ -2020,6 +2030,9 @@ def _inner_training_loop(
                         break
 
                     self.timers and self.timers("read-data").start()
+                    # Reset data loading timer for skipped data
+                    _data_load_time_for_global_step = 0.0
+                    _data_load_start_time = time.time()
                     continue
 
                 for inputs in inputs_list:
@@ -2189,7 +2202,6 @@ def hybrid_parallel_scale_param_grad(paramlist, hcg):
                         self.callback_handler.on_optimizer_begin(
                             args, self.state, self.control, scaler=self.scaler if self.do_grad_scaling else None
                         )
-
                         self.optimizer_step(args, model=model, parameters_list=parameters_list)
 
                         self.timers and self.timers("optimizer-step").stop()
@@ -2217,7 +2229,15 @@ def hybrid_parallel_scale_param_grad(paramlist, hcg):
 
                         self.control = self.callback_handler.on_step_end(args, self.state, self.control)
                         self._maybe_log_save_evaluate(tr_loss, model, epoch, ignore_keys_for_eval, inputs=inputs)
+                        # Log data loading time for this global_step
+                        logger.info(
+                            f"[DataLoad global_step: {self.state.global_step}] "
+                            f"data_load_time: {_data_load_time_for_global_step * 1000:.2f} ms "
+                            f"(accumulated over {args.gradient_accumulation_steps} micro-batches)"
+                        )
                         self._print_timer()
+                        # Reset data loading timer for next global_step
+                        _data_load_time_for_global_step = 0.0
                         step_control = 0
                     else:
                         self.control = self.callback_handler.on_substep_end(args, self.state, self.control)
@@ -2229,6 +2249,9 @@ def hybrid_parallel_scale_param_grad(paramlist, hcg):
                 if self.args.ignore_data_skip:
                     self.timers and self.timers("read-data").start()
 
+                # Reset start time for next iteration's data loading measurement
+                _data_load_start_time = time.time()
+
             if step < 0:
                 logger.warning(
                     f"There seems to be not a single sample in your epoch_iterator, stopping training at step"