Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions src/llmcompressor/args/dataset_arguments.py
Original file line number Diff line number Diff line change
Expand Up @@ -217,3 +217,6 @@ class DatasetArguments(CustomDatasetArguments):
"Default is set to True."
},
)

def is_dataset_provided(self) -> bool:
return self.dataset is not None or self.dataset_path is not None
2 changes: 1 addition & 1 deletion src/llmcompressor/entrypoints/oneshot.py
Original file line number Diff line number Diff line change
Expand Up @@ -125,7 +125,7 @@ def __init__(
self.output_dir = output_dir

# initialize the model and processor
pre_process(model_args)
pre_process(model_args, dataset_args, output_dir)

# Set instance attributes
self.model = self.model_args.model
Expand Down
4 changes: 2 additions & 2 deletions src/llmcompressor/entrypoints/train.py
Original file line number Diff line number Diff line change
Expand Up @@ -59,11 +59,11 @@ def train(**kwargs) -> PreTrainedModel:
```

"""
model_args, dataset_args, recipe_args, training_args, _ = parse_args(
model_args, dataset_args, recipe_args, training_args, output_dir = parse_args(
include_training_args=True, **kwargs
)

pre_process(model_args)
pre_process(model_args, dataset_args, output_dir)
dispatch_for_generation(model_args.model) # train is dispatched same as generation

processed_dataset = get_processed_dataset(
Expand Down
37 changes: 31 additions & 6 deletions src/llmcompressor/entrypoints/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,12 @@
)
from transformers.utils.quantization_config import CompressedTensorsConfig

from llmcompressor.args import ModelArguments, RecipeArguments, TrainingArguments
from llmcompressor.args import (
DatasetArguments,
ModelArguments,
RecipeArguments,
TrainingArguments,
)
from llmcompressor.core import reset_session
from llmcompressor.pytorch.model_load.helpers import parse_dtype
from llmcompressor.transformers.sparsification.compressed_tensors_utils import (
Expand All @@ -30,7 +35,11 @@
from llmcompressor.utils.fsdp.helpers import is_fsdp_model


def pre_process(model_args: "ModelArguments"):
def pre_process(
model_args: ModelArguments,
dataset_args: DatasetArguments,
output_dir: Optional[str],
):
"""
Prepares the model and tokenizer/processor for calibration.
- Initializes the model if it's specified as a path or string.
Expand All @@ -54,11 +63,27 @@ def pre_process(model_args: "ModelArguments"):
model_args.model = model
model_args.distill_teacher = distill_teacher

# Initialize processor
# Initialize processor if dataset provided
if isinstance(model_args.processor, (str, type(None))):
model_args.processor = initialize_processor_from_path(
model_args, model_args.model
)
try:
model_args.processor = initialize_processor_from_path(
model_args, model_args.model
)
except Exception as e:
if dataset_args.is_dataset_provided():
raise RuntimeError(
"An error occurred when attempting to initialize "
"model processor, which is required when a dataset "
"is provided. To resolve, create and pass in a "
"processor directly to `oneshot`/`train`."
) from e
elif output_dir:
logger.warning(
"Model processor could not be auto-initialized and "
"will not be saved along with the model. To resolve, "
"create and pass in a processor directly to "
f"`oneshot`/`train`.\nInitialization Error: {e}"
)

# untie tie_word_embeddings weights
if not model_args.tie_word_embeddings:
Expand Down
Loading