diff --git a/src/llmcompressor/args/model_arguments.py b/src/llmcompressor/args/model_arguments.py index 321a73a08..8cccc5cc0 100644 --- a/src/llmcompressor/args/model_arguments.py +++ b/src/llmcompressor/args/model_arguments.py @@ -6,7 +6,7 @@ class ModelArguments: """ Model variables used for post_train (post-training quantization, and sparsification), - finetuning and stage runners (sequential run of oneshot and finetune). + finetuning and stage runners (sequential run of post_train and finetune). """ diff --git a/src/llmcompressor/args/training_arguments.py b/src/llmcompressor/args/training_arguments.py index c4ebbbfea..8cdd77245 100644 --- a/src/llmcompressor/args/training_arguments.py +++ b/src/llmcompressor/args/training_arguments.py @@ -16,7 +16,7 @@ class TrainingArguments(HFTrainingArgs): """ - do_oneshot: Optional[bool] = field( + do_post_train: Optional[bool] = field( default=False, metadata={"help": "Whether to run one-shot calibration in stages"}, ) diff --git a/src/llmcompressor/transformers/finetune/data/data_helpers.py b/src/llmcompressor/transformers/finetune/data/data_helpers.py index 7d4e59068..ba43f95a6 100644 --- a/src/llmcompressor/transformers/finetune/data/data_helpers.py +++ b/src/llmcompressor/transformers/finetune/data/data_helpers.py @@ -98,7 +98,7 @@ def make_dataset_splits( do_train: bool = False, do_eval: bool = False, do_predict: bool = False, - do_oneshot: bool = False, + do_post_train: bool = False, ) -> Dict[str, Dataset]: """ Restructures the datasets dictionary based on what tasks will be run @@ -108,7 +108,7 @@ def make_dataset_splits( :param do_train: Whether to store the train dataset :param do_eval: Whether to store the validation dataset :param do_predict: Whether to store the test dataset - :param do_oneshot: Whether to store the calibration dataset + :param do_post_train: Whether to store the calibration dataset :return: Datasets to be used by the requested tasks """ @@ -132,11 +132,11 @@ def make_dataset_splits( if "test" not in tokenized_datasets: raise ValueError("--do_predict requires a test dataset") predict_split = tokenized_datasets["test"] - if do_oneshot: + if do_post_train: calib_split = tokenized_datasets.get("calibration") if calib_split is None: if "train" not in tokenized_datasets: - raise ValueError("--do_oneshot requires a calibration dataset") + raise ValueError("--do_post_train requires a calibration dataset") calib_split = tokenized_datasets["train"] split_datasets = { @@ -305,7 +305,7 @@ def _get_split_name(inp_str): datasets = make_dataset_splits( tokenized_datasets, - do_oneshot=True, + do_post_train=True, ) calibration_dataset = datasets.get("calibration") diff --git a/src/llmcompressor/transformers/finetune/runner.py b/src/llmcompressor/transformers/finetune/runner.py index 769b84248..dc563847a 100644 --- a/src/llmcompressor/transformers/finetune/runner.py +++ b/src/llmcompressor/transformers/finetune/runner.py @@ -124,7 +124,7 @@ def _get_split_name(inp_str): do_train=self._training_args.do_train, do_eval=self._training_args.do_eval, do_predict=self._training_args.do_predict, - do_oneshot=self._training_args.do_oneshot, + do_post_train=self._training_args.do_post_train, ) def get_dataset_split(self, split_name: str) -> Dataset: diff --git a/src/llmcompressor/transformers/finetune/text_generation.py b/src/llmcompressor/transformers/finetune/text_generation.py index 091fc5145..a198adac0 100644 --- a/src/llmcompressor/transformers/finetune/text_generation.py +++ b/src/llmcompressor/transformers/finetune/text_generation.py @@ -89,13 +89,13 @@ def post_train(**kwargs): """ # TODO: Get rid of training args when Oneshot refactor comes in model_args, data_args, recipe_args, training_args = parse_args(**kwargs) - training_args.do_oneshot = True + training_args.do_post_train = True main(model_args, data_args, recipe_args, training_args) # alias -def post_train**kwargs): +def oneshot(**kwargs): logger.warning( ("oneshot is now deprecated. Please use " "`post_train` method instead.") ) @@ -216,7 +216,7 @@ def initialize_model_from_path( # if running oneshot outside of FSDP, apply user device settings device_map = None fsdp_enabled = os.environ.get("ACCELERATE_USE_FSDP", "false") == "true" - if not fsdp_enabled and training_args.do_oneshot: + if not fsdp_enabled and training_args.do_post_train: device_map = training_args.post_train_device logger.warning(f"Moving {model_path} to device {device_map} for One-Shot") elif not fsdp_enabled: @@ -344,7 +344,7 @@ def main( for stage in recipe_obj.stages: run_type = stage.infer_run_type() if run_type is StageRunType.ONESHOT: - training_args.do_oneshot = True + training_args.do_post_train = True elif run_type is StageRunType.TRAIN: training_args.do_train = True @@ -438,7 +438,7 @@ def main( stage_runner.train(checkpoint) # One Shot - if training_args.do_oneshot: + if training_args.do_post_train: stage_runner.one_shot() # Evaluation diff --git a/tests/llmcompressor/transformers/finetune/data/test_dataset_loading.py b/tests/llmcompressor/transformers/finetune/data/test_dataset_loading.py index dcc602877..ed62e9ac9 100644 --- a/tests/llmcompressor/transformers/finetune/data/test_dataset_loading.py +++ b/tests/llmcompressor/transformers/finetune/data/test_dataset_loading.py @@ -324,7 +324,7 @@ def preprocess(sample): data_args=DatasetArguments( dataset=tokenized_dataset, shuffle_calibration_samples=False ), - training_args=TrainingArguments(do_oneshot=True), + training_args=TrainingArguments(do_post_train=True), recipe_args=RecipeArguments(), ) stage_runner.populate_datasets(processor=None)