Skip to content

Commit 4af2d25

Browse files
authored
[Callbacks][Docs] Add docstrings to saving functions (#1201)
## Purpose ## * Code documentation ## Prerequisites ## * #1160 --------- Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
1 parent c8af27d commit 4af2d25

File tree

2 files changed

+25
-4
lines changed

2 files changed

+25
-4
lines changed

src/llmcompressor/pytorch/model_load/helpers.py

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,11 +27,20 @@
2727
def save_checkpoint(
2828
save_path: str,
2929
model: PreTrainedModel,
30-
processor: Processor,
30+
processor: Optional[Processor] = None,
3131
save_safetensors: bool = True,
3232
save_compressed: bool = True,
3333
skip_sparsity_compression_stats: bool = False,
3434
):
35+
"""
36+
Save a model, processor, and recipe
37+
38+
:param save_path: Path used to save model and processor
39+
:param model: model to save
40+
:param processor: processor to save
41+
:param save_safetensors: save model checkpoint using safetensors file type
42+
:param save_compressed: save model checkpoint using compressed-tensors format
43+
"""
3544
# saving the model also saves the recipe
3645
model.save_pretrained(
3746
save_path,

src/llmcompressor/transformers/sparsification/compressed_tensors_utils.py

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,12 @@
3636
def modify_save_pretrained(model: PreTrainedModel):
3737
"""
3838
Overrides a PreTrainedModel's save_pretrained() method with a wrapped version that
39-
supports compression
39+
supports compression. The new save_pretrained function performs the following saving
40+
operations:
41+
42+
1. Saves the model state, potentially in a compressed format
43+
2. Saves the recipe, appending any current recipes to existing recipe files
44+
3. Copies any necessary python files from the model cache
4045
"""
4146

4247
def save_pretrained_compressed(save_pretrained_method):
@@ -279,9 +284,16 @@ def get_model_compressor(
279284
)
280285

281286

282-
def update_and_save_recipe(model_path: str, save_directory: str):
287+
def update_and_save_recipe(model_stub: str, save_directory: str):
288+
"""
289+
Save a recipe ontop of any existing recipe files located at model_stub
290+
291+
:param model_stub: path to existing model or model stub which may contain an
292+
existing recipe
293+
:param save_directory: path to save combined existing recipe and current recipe
294+
"""
283295
recipes_to_save = []
284-
existing_recipe = infer_recipe_from_model_path(model_path)
296+
existing_recipe = infer_recipe_from_model_path(model_stub)
285297
if existing_recipe is not None:
286298
recipes_to_save.append(existing_recipe)
287299

0 commit comments

Comments
 (0)