[StageRemoval] Remove Predict pathway (#1146)

SUMMARY:
* Remove predict pathway from `main`
* Remove predict pathway from `StageRunner`
* Remove logic to make predict dataset split
* Remove `max_predict_samples` from `DatasetArguments`
* Remove any docs/comment that has `predict` inside
* Rename `predicted_ids` to `output_ids`

TEST PLAN:
* Pass existing tests 
* Delete tests involving predict

---------

Signed-off-by: George <[email protected]>

Author: horheynm

examples/quantization_w8a8_fp8/whisper_example.py

@@ -35,8 +35,8 @@
     sample["array"], sampling_rate=sample["sampling_rate"], return_tensors="pt"
 ).input_features
 input_features = input_features.to(model.device)
-predicted_ids = model.generate(input_features, language="en", forced_decoder_ids=None)
-print(processor.batch_decode(predicted_ids, skip_special_tokens=False)[0])
+output_ids = model.generate(input_features, language="en", forced_decoder_ids=None)
+print(processor.batch_decode(output_ids, skip_special_tokens=False)[0])
 # Mr. Quilter is the apostle of the middle classes and we are glad to welcome his gospel
 print("==========================================")
 

src/llmcompressor/args/dataset_arguments.py

@@ -150,15 +150,6 @@ class DatasetArguments(CustomDatasetArguments):
             "of training examples to this value if set."
         },
     )
-    max_predict_samples: Optional[int] = field(
-        default=None,
-        metadata={
-            "help": (
-                "For debugging purposes or quicker training, truncate the number of "
-                "prediction examples to this value if set."
-            ),
-        },
-    )
     min_tokens_per_module: Optional[float] = field(
         default=None,
         metadata={

src/llmcompressor/args/training_arguments.py

@@ -26,8 +26,8 @@ class TrainingArguments(HFTrainingArgs):
     output_dir: str = field(
         default="./output",
         metadata={
-            "help": "The output directory where the model predictions and "
-            "checkpoints will be written."
+            "help": "The output directory where the model safetensors, "
+            "recipe, config, and optionally checkpoints will be written."
         },
     )
 

src/llmcompressor/transformers/finetune/data/data_helpers.py

@@ -97,17 +97,15 @@ def get_raw_dataset(
 def make_dataset_splits(
     tokenized_datasets: Dict[str, Any],
     do_train: bool = False,
-    do_predict: bool = False,
     do_oneshot: bool = False,
 ) -> Dict[str, Dataset]:
     """
     Restructures the datasets dictionary based on what tasks will be run
-    (train, predict)
+    train
 
     :param tokenized_datasets: dictionary of processed datasets
-    :param do_train: Whether to store the train dataset
-    :param do_predict: Whether to store the test dataset
     :param do_oneshot: Whether to store the calibration dataset
+
     :return: Datasets to be used by the requested tasks
     """
 
@@ -117,16 +115,12 @@ def make_dataset_splits(
         if isinstance(tokenized_datasets, Dataset):
             tokenized_datasets = {"train": tokenized_datasets}
 
-    train_split = predict_split = calib_split = None
+    train_split = calib_split = None
 
     if do_train:
         if "train" not in tokenized_datasets:
             raise ValueError("--do_train requires a train dataset")
         train_split = tokenized_datasets["train"]
-    if do_predict:
-        if "test" not in tokenized_datasets:
-            raise ValueError("--do_predict requires a test dataset")
-        predict_split = tokenized_datasets["test"]
     if do_oneshot:
         calib_split = tokenized_datasets.get("calibration")
         if calib_split is None:
@@ -136,7 +130,6 @@ def make_dataset_splits(
 
     split_datasets = {
         "train": train_split,
-        "test": predict_split,
         "calibration": calib_split,
     }
     return split_datasets

src/llmcompressor/transformers/finetune/runner.py

@@ -38,7 +38,7 @@ class StageRunner:
     LifeCycle
         - populate_datasets()
         - set_trainer()
-        - train() / predict()
+        - train()
 
     :param model_args: Arguments pertaining to model/config/processor
     :param data_args: Arguments pertaining to what data to use for different flows
@@ -121,7 +121,6 @@ def _get_split_name(inp_str):
         self.datasets = make_dataset_splits(
             tokenized_datasets,
             do_train=self._training_args.do_train,
-            do_predict=self._training_args.do_predict,
             do_oneshot=self._training_args.do_oneshot,
         )
 
@@ -155,18 +154,6 @@ def train(self, checkpoint: str, stage: Optional[str] = None):
         # this includes saving the state, optimizer and scheduler
         self.trainer.save_model(output_dir=self._output_dir)
 
-    def predict(self):
-        """
-        Run trainer's prediction loop on predict_dataset, logging the desired metrics
-        """
-        logger.info("*** Predict ***")
-        results = self.trainer.predict(self.dataset["test"])
-        metrics = results.metrics
-
-        metrics["predict_samples"] = len(self.dataset["test"])
-        self.trainer.log_metrics("predict", metrics)
-        self.trainer.save_metrics("predict", metrics)
-
     def run_sequential_stages(self, checkpoint: Optional[str] = None):
         """
         Run the recipe stage by stage, allowing for alternating between one-shot and

src/llmcompressor/transformers/finetune/session_mixin.py

@@ -344,31 +344,6 @@ def compute_loss(
 
         return loss
 
-    def prediction_step(
-        self,
-        model: Module,
-        inputs: Dict[str, Union[torch.Tensor, Any]],
-        prediction_loss_only: bool,
-        ignore_keys: Optional[List[str]] = None,
-    ) -> Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]:
-        """
-        Wraps the prediction step from the original trainer to remove any input entry
-        that should not be passed to the model.
-        This situation may arise when distillation is used and the teacher model
-        contains more inputs than the student model.
-        """
-        self._check_super_defined("prediction_step")
-
-        inputs = {k: inputs[k] for k in inputs if k in self._model_signature_columns}
-
-        model_outputs = super().prediction_step(
-            model=model,
-            inputs=inputs,
-            prediction_loss_only=prediction_loss_only,
-            ignore_keys=ignore_keys,
-        )
-        return model_outputs
-
     def train(self, *args, stage: Optional[str] = None, **kwargs):
         """
         Run a sparsification training cycle. Runs initialization for the sparse session
@@ -408,22 +383,6 @@ def train(self, *args, stage: Optional[str] = None, **kwargs):
 
         return output
 
-    def predict(self, *args, **kwargs):
-        """
-        Run a sparsification prediction cycle.
-        Runs initialize_structure for the sparse session before calling
-        super().predict() and finalization of the session after.
-
-        :param args: positional args to pass to super().predict()
-        :param kwargs: keyword args to pass to super().predict()
-        :return: the output from super.predict()
-        """
-        self.initialize_structure()
-        output = super().predict(*args, **kwargs)
-        self.finalize_session()
-
-        return output
-
     def save_model(self, output_dir: str, _internal_call=False, _is_oneshot=False):
         """
         Override of the save_model function and expects it to exist in the parent.

src/llmcompressor/transformers/finetune/text_generation.py

@@ -92,12 +92,13 @@ def eval(**kwargs):
 )
 def oneshot(**kwargs) -> None:
     from llmcompressor import oneshot
+
     oneshot(**kwargs)
 
 
 def apply(**kwargs):
     """
-    CLI entrypoint for any of training, predict or oneshot
+    CLI entrypoint for any of training, oneshot
     """
     report_to = kwargs.get("report_to", None)
     model_args, data_args, recipe_args, training_args = parse_args(**kwargs)
@@ -322,7 +323,8 @@ def main(
         - Trainer()
             - SessionMixIn()
             - HFTransformersTrainer()
-        - StageRunner.train() and/or predict() and/or oneshot()
+        - StageRunner.train() and/or  oneshot()
+
 
     :param model_args: Arguments pertaining to which model/config/tokenizer we are
     going to fine-tune from
@@ -437,10 +439,6 @@ def main(
             checkpoint = last_checkpoint
         stage_runner.train(checkpoint)
 
-    # Prediction
-    if training_args.do_predict:
-        stage_runner.predict()
-
     # save if model was provided as a string or custom output_dir was set
 
     if isinstance(model_args.model, str) or (

tests/llmcompressor/transformers/finetune/data/test_dataset_helpers.py

@@ -14,14 +14,11 @@ def test_combined_datasets():
     )
     raw_wikitext2 = get_raw_dataset(data_args)
     datasets = {"all": raw_wikitext2}
-
-    split_datasets = make_dataset_splits(datasets, do_train=True, do_predict=True)
+    split_datasets = make_dataset_splits(datasets, do_train=True)
     assert split_datasets.get("train") is not None
-    assert split_datasets.get("test") is not None
 
-    split_datasets = make_dataset_splits(datasets, do_train=True, do_predict=True)
+    split_datasets = make_dataset_splits(datasets, do_train=True)
     assert split_datasets.get("train") is not None
-    assert split_datasets.get("test") is not None
 
 
 @pytest.mark.unit
@@ -35,10 +32,11 @@ def test_separate_datasets():
         raw_wikitext2 = get_raw_dataset(data_args, split=split_str)
         datasets[split_name] = raw_wikitext2
 
-    split_datasets = make_dataset_splits(datasets, do_train=True, do_predict=False)
+    split_datasets = make_dataset_splits(datasets, do_train=True)
     assert split_datasets.get("train") is not None
-    assert split_datasets.get("test") is None
 
     with pytest.raises(ValueError):
         # fails due to no test split specified
-        split_datasets = make_dataset_splits(datasets, do_train=True, do_predict=True)
+
+        datasets.pop("train")
+        split_datasets = make_dataset_splits(datasets, do_train=True)