[Oneshot Refactor] dataclass Arguments

examples/trl_mixin/ex_trl_constant.py

@@ -3,7 +3,7 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from trl import DataCollatorForCompletionOnlyLM
 
-from llmcompressor.transformers import TrainingArguments
+from llmcompressor.transformers.utils.arg_parser import TrainingArguments
 
 model_path = "neuralmagic/Llama-2-7b-pruned50-retrained"
 output_dir = "./output_trl_sft_test_7b_gsm8k_sft_data"

examples/trl_mixin/ex_trl_distillation.py

@@ -1,9 +1,9 @@
 from sft_trainer import SFTTrainer
 from transformers import AutoModelForCausalLM, AutoTokenizer, DefaultDataCollator
 
-from llmcompressor.transformers import (
-    DataTrainingArguments,
-    TextGenerationDataset,
+from llmcompressor.transformers import TextGenerationDataset
+from llmcompressor.transformers.utils.arg_parser import (
+    DatasetArguments,
     TrainingArguments,
 )
 
@@ -21,7 +21,7 @@
 tokenizer = AutoTokenizer.from_pretrained(model_path)
 
 # Load gsm8k using SparseML dataset tools
-data_args = DataTrainingArguments(
+data_args = DatasetArguments(
     dataset="gsm8k", dataset_config_name="main", max_seq_length=512
 )
 dataset_manager = TextGenerationDataset.load_from_registry(

examples/trl_mixin/sft_trainer.py

@@ -1,7 +1,7 @@
 from trl import SFTConfig as TRLSFTConfig
 from trl import SFTTrainer as TRLSFTTrainer
 
-from llmcompressor.transformers import TrainingArguments
+from llmcompressor.transformers.utils.arg_parser import TrainingArguments
 from llmcompressor.transformers.finetune.session_mixin import SessionManagerMixIn
 
 __all__ = ["SFTTrainer"]

src/llmcompressor/transformers/finetune/README.md

@@ -74,9 +74,10 @@ train(
 
 Finetuning arguments are split up into 3 groups:
 
-* ModelArguments: `src/llmcompressor/transformers/finetune/model_args.py`
-* TrainingArguments: `src/llmcompressor/transformers/finetune/training_args.py`
-* DataTrainingArguments: `src/llmcompressor/transformers/finetune/data/data_training_args.py`
+* ModelArguments: `src/llmcompressor/transformers/utils/arg_parser/model_arguments.py`
+* TrainingArguments: `src/llmcompressor/transformers/utils/arg_parser/training_arguments.py`
+* DatasetArguments: `src/llmcompressor/transformers/utils/arg_parser/data_arguments.py`
+* RecipeArguments: `src/llmcompressor/transformers/utils/arg_parser/recipe_arguments.py`
 
 
 ## Running One-Shot with FSDP

src/llmcompressor/transformers/finetune/__init__.py

@@ -1,7 +1,5 @@
 # flake8: noqa
 
-from .data import DataTrainingArguments, TextGenerationDataset
-from .model_args import ModelArguments
+from .data import TextGenerationDataset
 from .session_mixin import SessionManagerMixIn
 from .text_generation import apply, compress, eval, oneshot, train
-from .training_args import TrainingArguments

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

@@ -4,7 +4,6 @@
 from .c4 import C4Dataset
 from .cnn_dailymail import CNNDailyMailDataset
 from .custom import CustomDataset
-from .data_args import DataTrainingArguments
 from .evolcodealpaca import EvolCodeAlpacaDataset
 from .flickr_30k import Flickr30K
 from .gsm8k import GSM8KDataset

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

@@ -8,12 +8,12 @@
 from datasets.formatting.formatting import LazyRow
 from loguru import logger
 
-from llmcompressor.transformers.finetune.data.data_args import DataTrainingArguments
 from llmcompressor.transformers.finetune.data.data_helpers import (
     LABELS_MASK_VALUE,
     get_custom_datasets_from_path,
     get_raw_dataset,
 )
+from llmcompressor.transformers.utils.arg_parser import DatasetArguments
 from llmcompressor.transformers.utils.preprocessing_functions import (
     PreprocessingFunctionRegistry,
 )
@@ -41,7 +41,7 @@ class TextGenerationDataset(RegistryMixin):
 
     def __init__(
         self,
-        data_args: DataTrainingArguments,
+        data_args: DatasetArguments,
         split: str,
         processor: Processor,
     ):

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

@@ -5,7 +5,7 @@
 from llmcompressor.typing import Processor
 
 if TYPE_CHECKING:
-    from llmcompressor.transformers import DataTrainingArguments as DataArgs
+    from llmcompressor.transformers.utils.arg_parser import DatasetArguments
 
 
 @TextGenerationDataset.register(name="c4")
@@ -18,7 +18,7 @@ class C4Dataset(TextGenerationDataset):
     :param processor: processor or tokenizer to use on dataset
     """
 
-    def __init__(self, data_args: "DataArgs", split: str, processor: Processor):
+    def __init__(self, data_args: "DatasetArguments", split: str, processor: Processor):
         data_args = deepcopy(data_args)
         data_args.dataset = "allenai/c4"
         data_args.text_column = "text"

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

@@ -5,7 +5,7 @@
 from llmcompressor.typing import Processor
 
 if TYPE_CHECKING:
-    from llmcompressor.transformers import DataTrainingArguments as DataArgs
+    from llmcompressor.transformers.utils.arg_parser import DatasetArguments
 
 
 @TextGenerationDataset.register(name="cnn_dailymail")
@@ -20,7 +20,7 @@ class CNNDailyMailDataset(TextGenerationDataset):
 
     SAMPLE_TEMPLATE = "Article:\n{article}\n\n### Summarization:\n{highlights}\n"
 
-    def __init__(self, data_args: "DataArgs", split: str, processor: Processor):
+    def __init__(self, data_args: "DatasetArguments", split: str, processor: Processor):
         data_args = deepcopy(data_args)
         data_args.dataset = "cnn_dailymail"
         data_args.dataset_config_name = "3.0.0"

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

@@ -5,7 +5,7 @@
 from llmcompressor.typing import Processor
 
 if TYPE_CHECKING:
-    from llmcompressor.transformers import DataTrainingArguments as DataArgs
+    from llmcompressor.transformers.utils.arg_parser import DatasetArguments
 
 
 @TextGenerationDataset.register(name="evolcodealpaca")
@@ -25,7 +25,7 @@ class EvolCodeAlpacaDataset(TextGenerationDataset):
         "\n\n### Response:\n"
     )
 
-    def __init__(self, data_args: "DataArgs", split: str, processor: Processor):
+    def __init__(self, data_args: "DatasetArguments", split: str, processor: Processor):
         data_args = deepcopy(data_args)
         data_args.dataset = "theblackcat102/evol-codealpaca-v1"
         data_args.text_column = "text"

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

@@ -7,7 +7,7 @@
 from llmcompressor.typing import Processor
 
 if TYPE_CHECKING:
-    from llmcompressor.transformers import DataTrainingArguments as DataArgs
+    from llmcompressor.transformers.utils.arg_parser import DatasetArguments
 
 
 @TextGenerationDataset.register(name="flickr", alias="flickr30k")
@@ -31,7 +31,7 @@ class Flickr30K(TextGenerationDataset):
         "{{ '<|assistant|>' }}\n{% endif %}\n{% endfor %}"
     )
 
-    def __init__(self, data_args: "DataArgs", split: str, processor: Processor):
+    def __init__(self, data_args: "DatasetArguments", split: str, processor: Processor):
         data_args = deepcopy(data_args)
         data_args.dataset = "lmms-lab/flickr30k"
 

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

@@ -5,7 +5,7 @@
 from llmcompressor.typing import Processor
 
 if TYPE_CHECKING:
-    from llmcompressor.transformers import DataTrainingArguments as DataArgs
+    from llmcompressor.transformers.utils.arg_parser import DatasetArguments
 
 
 @TextGenerationDataset.register(name="gsm8k")
@@ -20,7 +20,7 @@ class GSM8KDataset(TextGenerationDataset):
 
     GSM_TEMPLATE = "Question: {question}\nAnswer:"
 
-    def __init__(self, data_args: "DataArgs", split: str, processor: Processor):
+    def __init__(self, data_args: "DatasetArguments", split: str, processor: Processor):
         data_args = deepcopy(data_args)
         data_args.dataset = "gsm8k"
         data_args.text_column = "text"

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

@@ -5,7 +5,7 @@
 from llmcompressor.typing import Processor
 
 if TYPE_CHECKING:
-    from llmcompressor.transformers import DataTrainingArguments as DataArgs
+    from llmcompressor.transformers.utils.arg_parser import DatasetArguments
 
 
 @TextGenerationDataset.register(name="open_platypus")
@@ -28,7 +28,7 @@ class OpenPlatypusDataset(TextGenerationDataset):
         "instruction}\n\n### Response:\n",
     }
 
-    def __init__(self, data_args: "DataArgs", split: str, processor: Processor):
+    def __init__(self, data_args: "DatasetArguments", split: str, processor: Processor):
         data_args = deepcopy(data_args)
         data_args.dataset = "garage-bAInd/Open-Platypus"
         data_args.text_column = "text"

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

@@ -5,7 +5,7 @@
 from llmcompressor.typing import Processor
 
 if TYPE_CHECKING:
-    from llmcompressor.transformers import DataTrainingArguments as DataArgs
+    from llmcompressor.transformers.utils.arg_parser import DatasetArguments
 
 
 @TextGenerationDataset.register(name="ptb")
@@ -18,7 +18,7 @@ class PtbDataset(TextGenerationDataset):
     :param processor: processor or tokenizer to use on dataset
     """
 
-    def __init__(self, data_args: "DataArgs", split: str, processor: Processor):
+    def __init__(self, data_args: "DatasetArguments", split: str, processor: Processor):
         data_args = deepcopy(data_args)
         data_args.dataset = "ptb_text_only"
         data_args.text_column = "sentence"

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

@@ -7,7 +7,7 @@
 from llmcompressor.typing import Processor
 
 if TYPE_CHECKING:
-    from llmcompressor.transformers import DataTrainingArguments as DataArgs
+    from llmcompressor.transformers.utils.arg_parser import DatasetArguments
 
 
 @TextGenerationDataset.register(name="ultrachat_200k")
@@ -33,7 +33,7 @@ class UltraChatDataset(TextGenerationDataset):
         "{{ '<|assistant|>' }}\n{% endif %}\n{% endfor %}"
     )
 
-    def __init__(self, data_args: "DataArgs", split: str, processor: Processor):
+    def __init__(self, data_args: "DatasetArguments", split: str, processor: Processor):
         data_args = deepcopy(data_args)
         data_args.dataset = "HuggingFaceH4/ultrachat_200k"
         data_args.text_column = "messages"

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

@@ -5,7 +5,7 @@
 from llmcompressor.typing import Processor
 
 if TYPE_CHECKING:
-    from llmcompressor.transformers import DataTrainingArguments as DataArgs
+    from llmcompressor.transformers.utils.arg_parser import DatasetArguments
 
 
 @TextGenerationDataset.register(name="wikitext")
@@ -18,7 +18,7 @@ class WikiTextDataset(TextGenerationDataset):
     :param processor: processor or tokenizer to use on dataset
     """
 
-    def __init__(self, data_args: "DataArgs", split: str, processor: Processor):
+    def __init__(self, data_args: "DatasetArguments", split: str, processor: Processor):
         data_args = deepcopy(data_args)
         data_args.dataset = "Salesforce/wikitext"
         data_args.text_column = "text"

src/llmcompressor/transformers/finetune/runner.py

@@ -16,13 +16,16 @@
 from llmcompressor.pytorch.utils import tensors_to_device
 from llmcompressor.recipe import Recipe, StageRunType
 from llmcompressor.transformers.finetune.data import TextGenerationDataset
-from llmcompressor.transformers.finetune.data.data_args import DataTrainingArguments
 from llmcompressor.transformers.finetune.data.data_helpers import (
     format_calibration_data,
     make_dataset_splits,
 )
-from llmcompressor.transformers.finetune.model_args import ModelArguments
-from llmcompressor.transformers.finetune.training_args import TrainingArguments
+from llmcompressor.transformers.utils.arg_parser import (
+    DatasetArguments,
+    ModelArguments,
+    RecipeArguments,
+    TrainingArguments,
+)
 from llmcompressor.typing import Processor
 from llmcompressor.utils.fsdp.helpers import is_fsdp_model, save_model_and_recipe
 
@@ -46,13 +49,15 @@ class StageRunner:
 
     def __init__(
         self,
-        data_args: "DataTrainingArguments",
+        data_args: "DatasetArguments",
         model_args: "ModelArguments",
         training_args: "TrainingArguments",
+        recipe_args: "RecipeArguments",
     ):
         self._data_args = data_args
         self._model_args = model_args
         self._training_args = training_args
+        self._recipe_args = recipe_args
 
         self.datasets = {}
         self.trainer = None
@@ -214,7 +219,7 @@ def run_sequential_stages(self, checkpoint: Optional[str] = None):
         :param checkpoint: optional checkpoint to pick up a stage from
         """
 
-        recipe_obj = Recipe.create_instance(self._training_args.recipe)
+        recipe_obj = Recipe.create_instance(self._recipe_args.recipe)
         with self.trainer.accelerator.main_process_first():
             checkpoint_dir = self._model_args.model
             completed_stages = get_completed_stages(checkpoint_dir)
@@ -256,16 +261,13 @@ def run_sequential_stages(self, checkpoint: Optional[str] = None):
                 self.train(checkpoint=checkpoint, stage=stage_name)
             checkpoint = None
 
-            if (
-                self._training_args.output_dir
-                != TrainingArguments.__dataclass_fields__["output_dir"].default
-            ):
+            if self._training_args.output_dir:
                 save_model_and_recipe(
                     model=self.trainer.model,
                     save_path=self._output_dir,
                     processor=self.processor,
                     save_safetensors=self._training_args.save_safetensors,
-                    save_compressed=self._training_args.save_compressed,
+                    save_compressed=self._model_args.save_compressed,
                 )
 
             # save stage to checkpoint dir

src/llmcompressor/transformers/finetune/session_mixin.py

@@ -36,8 +36,10 @@
 from llmcompressor.utils.pytorch import qat_active
 
 if TYPE_CHECKING:
-    from llmcompressor.transformers import DataTrainingArguments
-
+    from llmcompressor.transformers.utils.arg_parser import (
+        DatasetArguments,
+        ModelArguments,
+    )
 
 __all__ = [
     "SessionManagerMixIn",
@@ -68,12 +70,14 @@ def __init__(
         self,
         recipe: Optional[str] = None,
         recipe_args: Optional[Union[Dict[str, Any], str]] = None,
-        data_args: Optional["DataTrainingArguments"] = None,
+        data_args: Optional["DatasetArguments"] = None,
+        model_args: Optional["ModelArguments"] = None,
         teacher: Optional[Union[Module, str]] = None,
         **kwargs,
     ):
         self.recipe = recipe
         self.recipe_args = recipe_args
+        self.model_args = model_args
         self.teacher = teacher
 
         # parse training and metadata args
@@ -374,16 +378,16 @@ def train(self, *args, stage: Optional[str] = None, **kwargs):
         self.initialize_session(epoch=epoch, checkpoint=checkpoint, stage=stage)
 
         # do not save checkpoints as compressed
-        original_save_compressed = self.args.save_compressed
-        self.args.save_compressed = False
+        original_save_compressed = self.model_args.save_compressed
+        self.model_args.save_compressed = False
 
         # train with accelerator
         self.accelerator.wait_for_everyone()
         output = super().train(*args, **kwargs)
         self.accelerator.wait_for_everyone()
 
         # restore original setting for saving final model
-        self.args.save_compressed = original_save_compressed
+        self.model_args.save_compressed = original_save_compressed
 
         # lifecycle
         self.finalize_session()
@@ -433,7 +437,6 @@ def one_shot(
     ):
         """
         Run oneshot calibration on the active model
-
         :param stage: which stage of the recipe to run, or None to run whole recipe
         :param calib_data: dataloader of calibration data
         """
@@ -474,15 +477,15 @@ def save_model(self, output_dir: str, _internal_call=False, _is_oneshot=False):
         if not is_fsdp_model(self.model):
             self.model.save_pretrained(
                 output_dir,
-                save_compressed=self.args.save_compressed,
+                save_compressed=self.model_args.save_compressed,
                 safe_serialization=self.args.save_safetensors,
             )
         else:  # FSDP model
             save_pretrained_fsdp(
                 model=self.model,
                 accelerator=self.accelerator,
                 output_dir=output_dir,
-                save_compressed=self.args.save_compressed,
+                save_compressed=self.model_args.save_compressed,
                 save_safetensors=self.metadata.get("save_safetensors", False),
             )