Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Update QuantizationRecipe to use checkpointer.save_checkpoint #2257

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Update QuantizationRecipe to use checkpointer.save_checkpoint #2257

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
66b7f09
Update checkpoint handling in QuantizationRecipe to use checkpointer.…
Ankur-singh Jan 13, 2025
b42649e
fix safe serialization for other checkpointers
Ankur-singh Jan 23, 2025
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
25 changes: 8 additions & 17 deletions recipes/quantize.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,10 +3,8 @@
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
import os
import sys
import time
from pathlib import Path
from typing import Any, Dict

import torch
Expand Down Expand Up @@ -54,6 +52,12 @@ def __init__(self, cfg: DictConfig) -> None:

def load_checkpoint(self, checkpointer_cfg: DictConfig) -> Dict[str, Any]:
self._checkpointer = config.instantiate(checkpointer_cfg)
if hasattr(self._checkpointer, "_safe_serialization"):
logger.info(
"Setting safe_serialization to False. TorchAO quantization is compatible "
"only with HuggingFace's non-safetensor serialization and deserialization."
)
self._checkpointer._safe_serialization = False
checkpoint_dict = self._checkpointer.load_checkpoint()
return checkpoint_dict

Expand Down Expand Up @@ -95,21 +99,8 @@ def quantize(self, cfg: DictConfig):
logger.info(f"Memory used: {torch.cuda.max_memory_allocated() / 1e9:.02f} GB")

def save_checkpoint(self, cfg: DictConfig):
ckpt_dict = self._model.state_dict()
file_name = cfg.checkpointer.checkpoint_files[0].split(".")[0]

output_dir = Path(cfg.checkpointer.output_dir)
output_dir.mkdir(exist_ok=True)
checkpoint_file = Path.joinpath(
output_dir, f"{file_name}-{self._quantization_mode}".rstrip("-qat")
).with_suffix(".pt")

torch.save(ckpt_dict, checkpoint_file)
logger.info(
"Model checkpoint of size "
f"{os.path.getsize(checkpoint_file) / 1024**3:.2f} GiB "
f"saved to {checkpoint_file}"
)
ckpt_dict = {training.MODEL_KEY: self._model.state_dict()}
self._checkpointer.save_checkpoint(ckpt_dict, epoch=0)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This does clean the code quite a bit but now the files are saved with an awkward name in an epoch-0 folder. I suppose that is alright for now since this is meant to be an example, and we don't have a way of changing the checkpoint subfolder name. But maybe at some point in the future it would be nice to control.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Totally agreed, the file name is a bit awkward. However, since epoch is a required argument, we can't get around it. We should circle back to this sometime in future.



@config.parse
Expand Down
Loading