Skip to content

Commit 7f49448

Browse files
author
George Ohashi
committed
comments
1 parent 9e26589 commit 7f49448

File tree

6 files changed

+396
-0
lines changed

6 files changed

+396
-0
lines changed

‎src/llmcompressor/args/README.md‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
# Input arguments for `oneshot`, `train`, `eval` entrypoints
2+
3+
Parsers in `llm-compressor` define the input arguments required for various entry points, including `oneshot`, `train`, and `eval`.
4+
5+
Each entry point (e.g., oneshot) carries out its logic based on the provided input arguments, `model`, `recipe`, and `dataset`.
6+
7+
```python
8+
from llmcompressor.transformers import oneshot
9+
10+
model = ...
11+
recipe = ...
12+
dataset = ...
13+
oneshot(model=model, recipe=recipe, dataset=dataset)
14+
```
15+
16+
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:
17+
18+
```python
19+
oneshot(
20+
...,
21+
output_dir=...,
22+
)
23+
```
24+
25+
These input arguments can be overloaded into the function signature and will be parsed using Hugging Face's [argument parser](https://github.com/huggingface/transformers/blob/main/src/transformers/hf_argparser.py). The parsers define the acceptable inputs; therefore any arguments to be passed in must be defined.
26+
27+
`llm-compressor` uses four parsers, located in `llm_compressor/arg_parser`:
28+
* ModelArguments
29+
* DatasetArguments
30+
* RecipeArguments
31+
* TrainingArguments
32+
33+
34+
## ModelArguments
35+
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.
36+
37+
## DataArguments
38+
Manages data loading and preprocessing. The dataset argument can specify a Hugging Face dataset stub or a local dataset compatible with [`load_dataset`](https://github.com/huggingface/datasets/blob/3a4e74a9ace62ecd5c9cde7dcb6bcabd65cc7857/src/datasets/load.py#L1905). The preprocessing_func is a callable function that applies custom logic, such as formatting the data using a chat template.
39+
40+
## RecipeArguments
41+
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.
42+
43+
## TrainingArguments
44+
Specifies training parameters based on Hugging Face's [TrainingArguments class](https://github.com/huggingface/transformers/blob/main/src/transformers/training_args.py). These parameters include settings like learning rate (`learning_rate`), and the optimizer to use (`optim`).
45+

‎src/llmcompressor/args/__init__.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
# flake8: noqa
2+
3+
from .dataset_arguments import DatasetArguments
4+
from .model_arguments import ModelArguments
5+
from .recipe_arguments import RecipeArguments
6+
from .training_arguments import TrainingArguments

‎src/llmcompressor/args/dataset_arguments.py‎

Lines changed: 189 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,189 @@
1+
from dataclasses import dataclass, field
2+
from typing import Any, Callable, Dict, List, Optional, Union
3+
4+
from transformers import DefaultDataCollator
5+
6+
7+
@dataclass
8+
class DVCDatasetArguments:
9+
"""
10+
Arguments for training using DVC
11+
"""
12+
13+
dvc_data_repository: Optional[str] = field(
14+
default=None,
15+
metadata={"help": "Path to repository used for dvc_dataset_path"},
16+
)
17+
18+
19+
@dataclass
20+
class CustomDatasetArguments(DVCDatasetArguments):
21+
"""
22+
Arguments for training using custom datasets
23+
"""
24+
25+
dataset_path: Optional[str] = field(
26+
default=None,
27+
metadata={
28+
"help": (
29+
"Path to the custom dataset. Supports json, csv, dvc. "
30+
"For DVC, the to dvc dataset to load, of format dvc://path. "
31+
"For csv or json, the path containing the dataset. "
32+
),
33+
},
34+
)
35+
36+
text_column: str = field(
37+
default="text",
38+
metadata={
39+
"help": (
40+
"Optional key to be used as the `text` input to tokenizer/processor "
41+
"after dataset preprocesssing"
42+
)
43+
},
44+
)
45+
46+
remove_columns: Union[None, str, List] = field(
47+
default=None,
48+
metadata={"help": "Column names to remove after preprocessing (deprecated)"},
49+
)
50+
51+
preprocessing_func: Union[None, str, Callable] = field(
52+
default=None,
53+
metadata={
54+
"help": (
55+
"Typically a function which applies a chat template. Can take the form "
56+
"of either a function to apply to the dataset, a name defined in "
57+
"src/llmcompressor/transformers/utils/preprocessing_functions.py, or "
58+
"a path to a function definition of the form /path/to/file.py:func"
59+
)
60+
},
61+
)
62+
63+
data_collator: Callable[[Any], Any] = field(
64+
default_factory=lambda: DefaultDataCollator(),
65+
metadata={"help": "The function to used to form a batch from the dataset"},
66+
)
67+
68+
69+
@dataclass
70+
class DatasetArguments(CustomDatasetArguments):
71+
"""
72+
Arguments pertaining to what data we are going to input our model for
73+
calibration, training or eval
74+
75+
Using `HfArgumentParser` we can turn this class into argparse
76+
arguments to be able to specify them on the command line
77+
"""
78+
79+
dataset: Optional[str] = field(
80+
default=None,
81+
metadata={
82+
"help": (
83+
"The name of the dataset to use (via the datasets library). "
84+
"Supports input as a string or DatasetDict from HF"
85+
)
86+
},
87+
)
88+
dataset_config_name: Optional[str] = field(
89+
default=None,
90+
metadata={
91+
"help": ("The configuration name of the dataset to use"),
92+
},
93+
)
94+
max_seq_length: int = field(
95+
default=384,
96+
metadata={
97+
"help": "The maximum total input sequence length after tokenization. "
98+
"Sequences longer than this will be truncated, sequences shorter will "
99+
"be padded."
100+
},
101+
)
102+
concatenate_data: bool = field(
103+
default=False,
104+
metadata={
105+
"help": "Whether or not to concatenate datapoints to fill max_seq_length"
106+
},
107+
)
108+
raw_kwargs: Dict = field(
109+
default_factory=dict,
110+
metadata={"help": "Additional keyboard args to pass to datasets load_data"},
111+
)
112+
splits: Union[None, str, List, Dict] = field(
113+
default=None,
114+
metadata={"help": "Optional percentages of each split to download"},
115+
)
116+
num_calibration_samples: Optional[int] = field(
117+
default=512,
118+
metadata={"help": "Number of samples to use for one-shot calibration"},
119+
)
120+
shuffle_calibration_samples: Optional[bool] = field(
121+
default=True,
122+
metadata={
123+
"help": "whether to shuffle the dataset before selecting calibration data"
124+
},
125+
)
126+
streaming: Optional[bool] = field(
127+
default=False,
128+
metadata={"help": "True to stream data from a cloud dataset"},
129+
)
130+
overwrite_cache: bool = field(
131+
default=False,
132+
metadata={"help": "Overwrite the cached preprocessed datasets or not."},
133+
)
134+
preprocessing_num_workers: Optional[int] = field(
135+
default=None,
136+
metadata={"help": "The number of processes to use for the preprocessing."},
137+
)
138+
pad_to_max_length: bool = field(
139+
default=True,
140+
metadata={
141+
"help": "Whether to pad all samples to `max_seq_length`. If False, "
142+
"will pad the samples dynamically when batching to the maximum length "
143+
"in the batch (which can be faster on GPU but will be slower on TPU)."
144+
},
145+
)
146+
max_train_samples: Optional[int] = field(
147+
default=None,
148+
metadata={
149+
"help": "For debugging purposes or quicker training, truncate the number "
150+
"of training examples to this value if set."
151+
},
152+
)
153+
max_eval_samples: Optional[int] = field(
154+
default=None,
155+
metadata={
156+
"help": "For debugging purposes or quicker training, truncate the number "
157+
"of evaluation examples to this value if set."
158+
},
159+
)
160+
max_predict_samples: Optional[int] = field(
161+
default=None,
162+
metadata={
163+
"help": (
164+
"For debugging purposes or quicker training, truncate the number of "
165+
"prediction examples to this value if set."
166+
),
167+
},
168+
)
169+
min_tokens_per_module: Optional[float] = field(
170+
default=None,
171+
metadata={
172+
"help": (
173+
"The minimum percentage of tokens (out of the total number) "
174+
"that the module should 'receive' throughout the forward "
175+
"pass of the calibration. If a module receives fewer tokens, "
176+
"a warning will be logged. Defaults to 1/num_of_experts."
177+
"note: this argument is only relevant for MoE models"
178+
),
179+
},
180+
)
181+
trust_remote_code_data: bool = field(
182+
default=False,
183+
metadata={
184+
"help": "Whether or not to allow for datasets defined on the Hub using "
185+
"a dataset script. This option should only be set to True for "
186+
"repositories you trust and in which you have read the code, as it "
187+
"will execute code present on the Hub on your local machine."
188+
},
189+
)

‎src/llmcompressor/args/model_arguments.py‎

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
from dataclasses import dataclass, field
2+
from typing import Optional
3+
4+
5+
@dataclass
6+
class ModelArguments:
7+
"""
8+
Model variables used for oneshot calibration, finetuning and
9+
stage runners (sequential run of oneshot and finetune).
10+
11+
"""
12+
13+
model: str = field(
14+
metadata={
15+
"help": (
16+
"A pretrained model or a string as a path to pretrained model, "
17+
"HF stub, or model identifier from huggingface.co/models."
18+
)
19+
},
20+
)
21+
distill_teacher: Optional[str] = field(
22+
default=None,
23+
metadata={
24+
"help": "Teacher model (a trained text generation model)",
25+
},
26+
)
27+
config_name: Optional[str] = field(
28+
default=None,
29+
metadata={
30+
"help": "Pretrained config name or path if not the same as model_name"
31+
},
32+
)
33+
tokenizer: Optional[str] = field(
34+
default=None,
35+
metadata={
36+
"help": "Pretrained tokenizer name or path if not the same as model_name"
37+
},
38+
)
39+
processor: Optional[str] = field(
40+
default=None,
41+
metadata={
42+
"help": "Pretrained processor name or path if not the same as model_name"
43+
},
44+
)
45+
cache_dir: Optional[str] = field(
46+
default=None,
47+
metadata={"help": "Where to store the pretrained data from huggingface.co"},
48+
)
49+
50+
use_auth_token: bool = field(
51+
default=False,
52+
metadata={
53+
"help": "Will use token generated when running `transformers-cli login` "
54+
"(necessary to use this script with private models)"
55+
},
56+
)
57+
precision: str = field(
58+
default="auto",
59+
metadata={"help": "Precision to cast model weights to, default to auto"},
60+
)
61+
62+
tie_word_embeddings: bool = field(
63+
default=False,
64+
metadata={
65+
"help": "Whether the model's input and output word embeddings "
66+
"should be tied. Note that this is only relevant if the "
67+
"model has a output word embedding layer."
68+
},
69+
)
70+
trust_remote_code_model: bool = field(
71+
default=False,
72+
metadata={
73+
"help": "Whether or not to allow for custom models to execute their "
74+
"own modeling files. This option should only be set to True for "
75+
"repositories you trust and in which you have read the code"
76+
},
77+
)
78+
save_compressed: Optional[bool] = field(
79+
default=True,
80+
metadata={"help": "Whether to compress sparse models during save"},
81+
)
82+
oneshot_device: Optional[str] = field(
83+
default="cuda:0",
84+
metadata={"help": "Device to run oneshot calibration on"},
85+
)
86+
model_revision: str = field(
87+
default="main",
88+
metadata={
89+
"help": "The specific model version to use "
90+
"(can be a branch name, tag name or commit id)"
91+
},
92+
)

‎src/llmcompressor/args/recipe_arguments.py‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
from dataclasses import dataclass, field
2+
from typing import List, Optional
3+
4+
5+
@dataclass
6+
class RecipeArguments:
7+
"""Recipe and session variables"""
8+
9+
recipe: Optional[str] = field(
10+
default=None,
11+
metadata={
12+
"help": "Path to a LLM Compressor sparsification recipe",
13+
},
14+
)
15+
recipe_args: Optional[List[str]] = field(
16+
default=None,
17+
metadata={
18+
"help": (
19+
"List of recipe arguments to evaluate, of the format key1=value1 "
20+
"key2=value2"
21+
)
22+
},
23+
)
24+
clear_sparse_session: Optional[bool] = field(
25+
default=False,
26+
metadata={
27+
"help": (
28+
"Whether to clear CompressionSession/CompressionLifecycle ",
29+
"data between runs.",
30+
)
31+
},
32+
)

0 commit comments

Comments
 (0)