Skip to content

README.md

Latest commit

 

History

History
45 lines (31 loc) · 2.34 KB

README.md

File metadata and controls

45 lines (31 loc) · 2.34 KB

Argument parser for oneshot, train, eval entrypoints

Argument parsers define the input arguments required by any entrypoint in llm-compressor. Supported entrypoint includes: oneshot, train or eval.

The logic is carried out with respect to the given input arguments, such as model, recipe and dataset.

from llmcompressor.transformers import oneshot

model = ...
recipe = ...
dataset = ...
oneshot(model=model, recipe=recipe, dataset=dataset)

In addition, users can futher control execution by providing additional arguments. For example, to save the optimized model after completion, the output_dir parameter can be specified:

oneshot(
    ..., 
    output_dir=...,
)

These input arguments can be overloaded into the function signature and will be parsed using Hugging Face's argument parser. The parsers define the acceptable inputs; therefore any arguments to be passed in must be defined.

llm-compressor uses four parsers, located in llm_compressor/arg_parser:

  • ModelArguments
  • DataArguments
  • RecipeArguments
  • TrainingArguments

ModelArguments

Handles model loading and saving. For example, ModelArguments.model can be a Hugging Face model identifier or an instance of AutoModelForCausalLM. The save_compressed flag is a boolean that determines whether the model is saved in compressed safetensors format to minimize disk usage.

DataArguments

Manages data loading and preprocessing. The dataset argument can specify a Hugging Face dataset stub or a local dataset compatible with load_dataset. The preprocessing_func is a callable function that applies custom logic, such as formatting the data using a chat template.

RecipeArguments

Defines the model recipe. A recipe consists of user-defined instructions for optimizing the model. Examples of recipes can be found in the /examples directory.

TrainingArguments

Specifies training parameters based on Hugging Face's TrainingArguments class. These parameters include settings like learning rate (learning_rate), and the optimizer to use (optim).