[Trainer] Support skip data intervals

paddlenlp/trainer/argparser.py

@@ -24,7 +24,17 @@
 from enum import Enum
 from inspect import isclass
 from pathlib import Path
-from typing import Any, Dict, Iterable, NewType, Optional, Tuple, Union, get_type_hints
+from typing import (
+    Any,
+    Dict,
+    Iterable,
+    NewType,
+    Optional,
+    Tuple,
+    Union,
+    get_args,
+    get_type_hints,
+)
 
 DataClass = NewType("DataClass", Any)
 DataClassType = NewType("DataClassType", Any)
@@ -129,7 +139,13 @@ def _parse_dataclass_field(parser: ArgumentParser, field: dataclasses.Field):
                 # This is the value that will get picked if we do --field_name (without value)
                 kwargs["const"] = True
         elif isclass(origin_type) and issubclass(origin_type, list):
-            kwargs["type"] = field.type.__args__[0]
+            # supprt one dimension list and two dimension list
+            if hasattr(get_args(field.type)[0], "__args__"):
+                kwargs["type"] = field.type.__args__[0].__args__[0]
+                kwargs["action"] = "append"
+            else:
+                kwargs["type"] = field.type.__args__[0]
+
             kwargs["nargs"] = "+"
             if field.default_factory is not dataclasses.MISSING:
                 kwargs["default"] = field.default_factory()

paddlenlp/trainer/trainer.py

@@ -138,6 +138,7 @@
     get_scheduler,
     has_length,
     set_seed,
+    should_skip_data,
     speed_metrics,
 )
 from .training_args import TrainingArguments
@@ -277,7 +278,7 @@
         # Seed must be set before instantiating the model when using model
         set_seed(seed=self.args.seed)
 
-        if model is None:
+        if model is None and not args.debug_data:
             raise RuntimeError("`Trainer` requires either a `model` or `model_init` argument")
 
         if self.args.to_static:
@@ -339,7 +340,7 @@
                 "You should subclass `Trainer` and override the `create_optimizer_and_scheduler` method."
             )
 
-        if self.args.pipeline_parallel_degree > 1 and self.args.use_hybrid_parallel:
+        if self.args.pipeline_parallel_degree > 1 and self.args.use_hybrid_parallel and not args.debug_data:
             from paddle.distributed.fleet.meta_parallel import PipelineLayer
 
             assert (isinstance(model, LoRAModel) and isinstance(model.model, PipelineLayer)) or isinstance(
@@ -357,6 +358,7 @@
         self._load_ckpt_func = dist.load_state_dict if self.args.enable_auto_parallel else paddle.load
         if self.args.use_async_save:
             self._async_optimizer_saver = AsyncSaver()
+        self.skip_global_steps = 0
 
         if args.max_steps > 0:
             logger.info("max_steps is given, it will override any value given in num_train_epochs")
@@ -377,26 +379,28 @@
 
         self.do_grad_scaling = False
         self.enable_autocast_context_manager = False
-        if args.fp16 or args.bf16:
-            # set do_grad_scaling, enable_autocast_context_manager
-            self._wrap_amp_model(args, model)
 
-        if args.recompute:
+        if not args.debug_data:
+            if args.fp16 or args.bf16:
+                # set do_grad_scaling, enable_autocast_context_manager
+                self._wrap_amp_model(args, model)
 
-            def fn(layer):
-                if hasattr(layer, "enable_recompute") and (
-                    layer.enable_recompute is False or layer.enable_recompute == 0
-                ):
-                    layer.enable_recompute = True
+            if args.recompute:
 
-            model.apply(fn)
+                def fn(layer):
+                    if hasattr(layer, "enable_recompute") and (
+                        layer.enable_recompute is False or layer.enable_recompute == 0
+                    ):
+                        layer.enable_recompute = True
 
-        default_label_names = (
-            ["start_positions", "end_positions"]
-            if "QusetionAnswering" in type(self.model).__name__ or "UIE" in type(self.model).__name__
-            else ["labels"]
-        )
-        self.label_names = default_label_names if self.args.label_names is None else self.args.label_names
+                model.apply(fn)
+
+            default_label_names = (
+                ["start_positions", "end_positions"]
+                if "QusetionAnswering" in type(self.model).__name__ or "UIE" in type(self.model).__name__
+                else ["labels"]
+            )
+            self.label_names = default_label_names if self.args.label_names is None else self.args.label_names
 
         self.control = self.callback_handler.on_init_end(self.args, self.state, self.control)
         self.print_config()
@@ -924,6 +928,7 @@
             step_control = 0  # used in loop control, reset to 0 after every step
             self.control = self.callback_handler.on_epoch_begin(args, self.state, self.control)
 
+            step = -1
             for step, inputs in enumerate(epoch_iterator):
                 if self.args.use_hybrid_parallel and self.args.sep_parallel_degree > 1:
                     inputs = split_inputs_sequence_dim(inputs)
@@ -960,6 +965,31 @@
                     steps_trained_progress_bar.close()
                     steps_trained_progress_bar = None
 
+                # Skip data
+                if should_skip_data(self.state.global_step, self.args.skip_data_intervals):
+                    logger.warning(f"Skip data at global step {self.state.global_step+1}, sub step {step_control}")
+                    logger.warning(f"{self.tokenizer.batch_decode(inputs['input_ids'], skip_special_tokens=True)}")
+
+                    if (step_control + 1) % args.gradient_accumulation_steps == 0 or (
+                        # last step in epoch but step is always smaller than gradient_accumulation_steps
+                        steps_in_epoch <= args.gradient_accumulation_steps
+                        and (step + 1) == steps_in_epoch
+                    ):
+                        self.skip_global_steps += 1
+                        self.state.global_step += 1
+                        self.state.epoch = epoch + (step + 1) / steps_in_epoch
+                        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)
+                        self._print_timer()
+                        step_control = 0
+                    else:
+                        self.control = self.callback_handler.on_substep_end(args, self.state, self.control)
+                        step_control += 1
+                    if self.control.should_epoch_stop or self.control.should_training_stop:
+                        break
+                    self.timers and self.timers("read-data").start()
+                    continue
+
                 if step_control % args.gradient_accumulation_steps == 0:
                     self.control = self.callback_handler.on_step_begin(args, self.state, self.control)
                     self.timers and self.timers("forward-backward").start()
@@ -1181,7 +1211,10 @@
                         )
 
         self._total_loss_scalar += tr_loss.item()
-        train_loss = self._total_loss_scalar / self.state.global_step
+        if self.state.global_step == self.skip_global_steps:
+            train_loss = 0.0
+        else:
+            train_loss = self._total_loss_scalar / (self.state.global_step - self.skip_global_steps)
 
         metrics = speed_metrics("train", start_time, num_samples=num_train_samples, num_steps=self.state.max_steps)
 

paddlenlp/trainer/trainer_utils.py

@@ -1100,3 +1100,20 @@
         tracker.add("global_seed", global_seed)
     if "local_seed" not in tracker.states_ and local_seed not in tracker.seeds_:
         tracker.add("local_seed", local_seed)
+
+
+def should_skip_data(global_step, skip_data_intervals):
+    """Whether to skip current step data"""
+
+    if skip_data_intervals is None:
+        return False
+    skip_flag = False
+    for interval in skip_data_intervals:
+        if len(interval) != 2 or interval[0] > interval[1] or interval[0] <= 0:
+            raise ValueError(f"Please check your skip interval {interval}")
+
+        start_global_step, end_step = interval[0], interval[1]
+        if start_global_step <= global_step + 1 <= end_step:
+            skip_flag = True
+            break
+    return skip_flag

paddlenlp/trainer/training_args.py

@@ -850,6 +850,16 @@ class TrainingArguments:
     release_grads: Optional[bool] = field(
         default=False, metadata={"help": "Whether to release gradients during training. Default is `False`."}
     )
+    skip_data_intervals: Optional[List[List[int]]] = field(
+        default=None,
+        metadata={"help": "The intervals to skip, pass start global step and end global step at each interval"},
+    )
+    debug_data: Optional[bool] = field(
+        default=False,
+        metadata={
+            "help": "whether to debug data.If set to True, will print the skip intervals data and skip training process."
+        },
+    )
 
     def __post_init__(self):
         env_local_rank = int(os.environ.get("PADDLE_RANK_IN_NODE", -1))