[Oneshot Refactor] dataclass Arguments

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(

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,20 @@
 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.transformers.utils.arg_parser.training_arguments import (
+    DEFAULT_OUTPUT_DIR,
+)
+from llmcompressor.transformers.utils.arg_parser.utils import get_dataclass_as_dict
 from llmcompressor.typing import Processor
 from llmcompressor.utils.fsdp.helpers import is_fsdp_model, save_model_and_recipe
 
@@ -46,13 +53,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 +223,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)
@@ -251,21 +260,30 @@ def run_sequential_stages(self, checkpoint: Optional[str] = None):
 
             # run stage
             if run_type is StageRunType.ONESHOT:
-                self.one_shot(stage=stage_name)
+                from llmcompressor.transformers.calibration import Oneshot
+
+                model = get_session_model()
+                self._model_args.model = model
+
+                oneshot = Oneshot(
+                    output_dir=self._training_args.output_dir,
+                    **get_dataclass_as_dict(self._model_args, ModelArguments),
+                    **get_dataclass_as_dict(self._data_args, DatasetArguments),
+                    **get_dataclass_as_dict(self._recipe_args, RecipeArguments),
+                )
+
+                oneshot.run(stage_name=stage_name)
             elif run_type is StageRunType.TRAIN:
                 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 != DEFAULT_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