[WIP] [Design] LLMCompressor Class

examples/multimodal_audio/whisper_example.py

@@ -1,15 +1,14 @@
 import torch
 from datasets import load_dataset
-from transformers import WhisperProcessor
+from transformers import WhisperForConditionalGeneration, WhisperProcessor
 
 from llmcompressor import oneshot
 from llmcompressor.modifiers.quantization import GPTQModifier
-from llmcompressor.transformers.tracing import TraceableWhisperForConditionalGeneration
 
 # Select model and load it.
 MODEL_ID = "openai/whisper-large-v3"
 
-model = TraceableWhisperForConditionalGeneration.from_pretrained(
+model = WhisperForConditionalGeneration.from_pretrained(
     MODEL_ID,
     device_map="auto",
     torch_dtype="auto",

src/llmcompressor/__init__.py

@@ -8,6 +8,8 @@
 
 # flake8: noqa
 
+from .core import LLMCompressor
+from .entrypoints import oneshot, post_train, train
 from .logger import LoggerConfig, configure_logger, logger
 from .version import (
     __version__,
@@ -32,12 +34,8 @@
     "configure_logger",
     "logger",
     "LoggerConfig",
+    "LLMCompressor",
+    "oneshot",
+    "post_train",
+    "train",
 ]
-
-from llmcompressor.core.session_functions import (
-    active_session,
-    callbacks,
-    create_session,
-    reset_session,
-)
-from llmcompressor.entrypoints import Oneshot, oneshot, train

src/llmcompressor/args/__init__.py

@@ -2,6 +2,5 @@
 
 from .dataset_arguments import DatasetArguments
 from .model_arguments import ModelArguments
-from .recipe_arguments import RecipeArguments
+from .post_train_arguments import PostTrainArguments
 from .training_arguments import TrainingArguments
-from .utils import parse_args

src/llmcompressor/args/dataset_arguments.py

@@ -3,71 +3,11 @@
 
 from transformers import DefaultDataCollator
 
-
-@dataclass
-class DVCDatasetArguments:
-    """
-    Arguments for training using DVC
-    """
-
-    dvc_data_repository: Optional[str] = field(
-        default=None,
-        metadata={"help": "Path to repository used for dvc_dataset_path"},
-    )
-
-
-@dataclass
-class CustomDatasetArguments(DVCDatasetArguments):
-    """
-    Arguments for training using custom datasets
-    """
-
-    dataset_path: Optional[str] = field(
-        default=None,
-        metadata={
-            "help": (
-                "Path to the custom dataset. Supports json, csv, dvc. "
-                "For DVC, the to dvc dataset to load, of format dvc://path. "
-                "For csv or json, the path containing the dataset. "
-            ),
-        },
-    )
-
-    text_column: str = field(
-        default="text",
-        metadata={
-            "help": (
-                "Optional key to be used as the `text` input to tokenizer/processor "
-                "after dataset preprocesssing"
-            )
-        },
-    )
-
-    remove_columns: Union[None, str, List] = field(
-        default=None,
-        metadata={"help": "Column names to remove after preprocessing (deprecated)"},
-    )
-
-    preprocessing_func: Union[None, str, Callable] = field(
-        default=None,
-        metadata={
-            "help": (
-                "Typically a function which applies a chat template. Can take the form "
-                "of either a function to apply to the dataset, a name defined in "
-                "src/llmcompressor/transformers/utils/preprocessing_functions.py, or "
-                "a path to a function definition of the form /path/to/file.py:func"
-            )
-        },
-    )
-
-    data_collator: Callable[[Any], Any] = field(
-        default_factory=lambda: DefaultDataCollator(),
-        metadata={"help": "The function to used to form a batch from the dataset"},
-    )
+from llmcompressor.typing import DatasetType
 
 
 @dataclass
-class DatasetArguments(CustomDatasetArguments):
+class DatasetArguments:
     """
     Arguments pertaining to what data we are going to input our model for
     calibration, training
@@ -76,7 +16,7 @@ class DatasetArguments(CustomDatasetArguments):
     arguments to be able to specify them on the command line
     """
 
-    dataset: Optional[str] = field(
+    dataset: Union[str, DatasetType, None] = field(
         default=None,
         metadata={
             "help": (
@@ -109,10 +49,7 @@ class DatasetArguments(CustomDatasetArguments):
         default_factory=dict,
         metadata={"help": "Additional keyboard args to pass to datasets load_data"},
     )
-    splits: Union[None, str, List, Dict] = field(
-        default=None,
-        metadata={"help": "Optional percentages of each split to download"},
-    )
+    split: Optional[str] = field(default=None)
     num_calibration_samples: Optional[int] = field(
         default=512,
         metadata={"help": "Number of samples to use for one-shot calibration"},
@@ -171,3 +108,50 @@ class DatasetArguments(CustomDatasetArguments):
             "will execute code present on the Hub on your local machine."
         },
     )
+
+    dataset_path: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": (
+                "Path to the custom dataset. Supports json, csv, dvc. "
+                "For DVC, the to dvc dataset to load, of format dvc://path. "
+                "For csv or json, the path containing the dataset. "
+            ),
+        },
+    )
+
+    text_column: str = field(
+        default="text",
+        metadata={
+            "help": (
+                "Optional key to be used as the `text` input to tokenizer/processor "
+                "after dataset preprocesssing"
+            )
+        },
+    )
+    remove_columns: Union[None, str, List] = field(
+        default=None,
+        metadata={"help": "Column names to remove after preprocessing (deprecated)"},
+    )
+
+    preprocessing_func: Union[None, str, Callable] = field(
+        default=None,
+        metadata={
+            "help": (
+                "Typically a function which applies a chat template. Can take the form "
+                "of either a function to apply to the dataset, a name defined in "
+                "src/llmcompressor/transformers/utils/preprocessing_functions.py, or "
+                "a path to a function definition of the form /path/to/file.py:func"
+            )
+        },
+    )
+
+    data_collator: Callable[[Any], Any] = field(
+        default_factory=lambda: DefaultDataCollator(),
+        metadata={"help": "The function to used to form a batch from the dataset"},
+    )
+
+    dvc_data_repository: Optional[str] = field(
+        default=None,
+        metadata={"help": "Path to repository used for dvc_dataset_path"},
+    )

src/llmcompressor/args/model_arguments.py

@@ -1,24 +1,26 @@
 from dataclasses import dataclass, field
 from typing import Optional
 
+from llmcompressor.typing import ModelInput, RecipeInput
+
 
 @dataclass
 class ModelArguments:
     """
     Model variables used for oneshot calibration, finetuning and
     stage runners (sequential run of oneshot and finetune).
-
     """
 
-    model: str = field(
+    model: ModelInput = field(
         metadata={
             "help": (
                 "A pretrained model or a string as a path to pretrained model, "
                 "HF stub, or model identifier from huggingface.co/models."
             )
         },
     )
-    distill_teacher: Optional[str] = field(
+    recipe: RecipeInput = field(metadata={"help": ""})
+    distill_teacher: ModelInput = field(
         default=None,
         metadata={
             "help": "Teacher model (a trained text generation model)",
@@ -79,10 +81,6 @@ class ModelArguments:
         default=True,
         metadata={"help": "Whether to compress sparse models during save"},
     )
-    oneshot_device: Optional[str] = field(
-        default="cuda:0",
-        metadata={"help": "Device to run oneshot calibration on"},
-    )
     model_revision: str = field(
         default="main",
         metadata={

src/llmcompressor/args/post_train_arguments.py

@@ -0,0 +1,16 @@
+from dataclasses import dataclass, field
+from typing import Optional
+
+import torch
+
+
+@dataclass
+class PostTrainArguments:
+    pipeline: Optional[str] = field(
+        default=None,
+        metadata={"help": "Calibration pipeline used to calibrate model."},
+    )
+
+    oneshot_device: Optional[torch.device] = field(default=None)
+
+    output_dir: Optional[str] = field(default=None)

src/llmcompressor/args/training_arguments.py

@@ -16,12 +16,9 @@ class TrainingArguments(HFTrainingArgs):
 
     """
 
-    do_oneshot: Optional[bool] = field(
-        default=False,
-        metadata={"help": "Whether to run one-shot calibration in stages"},
-    )
-    run_stages: Optional[bool] = field(
-        default=False, metadata={"help": "Whether to trigger recipe stage by stage"}
+    max_seq_length: Optional[int] = field(
+        default=512,
+        metadata={"help": "The maximum sequence length to use during training"},
     )
     output_dir: str = field(
         default="./output",