Update QuantizationRecipe to use checkpointer.save_checkpoint

[Ankur-singh]

Context

What is the purpose of this PR? Is it to

• add a new feature
• fix a bug
• update tests and/or documentation
• other (refactor)

Please link to any issues this PR addresses #2229

Changelog

What are the changes made in this PR?

• Simplified save_checkpoint method by making use of checkpointer.save_checkpoint instead of custom implementation.

Test plan

Please make sure to do each of the following if applicable to your PR. If you're unsure about any one of these just ask and we will happily help. We also have a contributing page for some guidance on contributing.

• run pre-commit hooks and linters (make sure you've first installed via pre-commit install)
• add unit tests for any new functionality
• update docstrings for any new or updated methods or classes
• run unit tests via pytest tests
• run recipe tests via pytest tests -m integration_test
• manually run any new or modified recipes with sufficient proof of correctness
• include relevant commands and any other artifacts in this summary (pastes of loss curves, eval results, etc.)

Quantization recipe:

# Config for QuantizationRecipe in quantize.py
#
# To launch, run the following command from root torchtune directory:
#    tune download Qwen/Qwen2.5-0.5B-Instruct --output-dir /tmp/Qwen2_5-0_5B-Instruct
#    tune run quantize --config quantization

output_dir: ./quantized # /tmp may be deleted by your system. Change it to your preference.

#
# Model arguments
model:
  _component_: torchtune.models.qwen2_5.qwen2_5_0_5b

checkpointer:
  _component_: torchtune.training.FullModelHFCheckpointer
  checkpoint_dir: /tmp/Qwen2_5-0_5B-Instruct
  checkpoint_files: [model.safetensors]
  recipe_checkpoint: null
  output_dir: ${output_dir}
  model_type: QWEN2

device: cuda
dtype: bf16
seed: 1234

quantizer:
  _component_: torchtune.training.quantization.Int8DynActInt4WeightQuantizer
  groupsize: 256

Output:

(tune) ➜  torchtune git:(refactor/quantization-recipe-save-checkpoint) ✗ tune run quantize --config ./custom_quant.yaml
Running QuantizationRecipe with resolved config:

checkpointer:
  _component_: torchtune.training.FullModelHFCheckpointer
  checkpoint_dir: /tmp/Qwen2_5-0_5B-Instruct
  checkpoint_files:
  - model.safetensors
  model_type: QWEN2
  output_dir: ./quantized
  recipe_checkpoint: null
device: cuda
dtype: bf16
model:
  _component_: torchtune.models.qwen2_5.qwen2_5_0_5b
output_dir: ./quantized
quantizer:
  _component_: torchtune.training.quantization.Int8DynActInt4WeightQuantizer
  groupsize: 256
seed: 1234

Setting manual seed to local seed 1234. Local seed is seed + rank = 1234 + 0
Model is initialized with precision torch.bfloat16.
Time for quantization: 0.08 sec
Memory used: 1.45 GB
Model checkpoint of size 0.82 GiB saved to quantized/epoch_0/ft-model-00001-of-00001.bin
Saving final epoch checkpoint.
The full model checkpoint, including all weights and configurations, has been saved successfully.You can now use this checkpoint for further training or inference.

UX

If your function changed a public API, please add a dummy example of what the user experience will look like when calling it.
Here is a docstring example
and a tutorial example

• I did not change any public API
• I have added an example to docs or docstrings

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/2257

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[Ankur-singh]

@felipemello1 There are two caveats that I think are important to highlight

• epoch is set to 0 as it is a required parameter. Is there a better and elegant way?

• safe_serialization is disabled because

  torchao quantization is implemented with tensor subclasses, it only work with huggingface non-safetensor serialization and deserialization

  source

[felipemello1]

hey @Ankur-singh , thanks for the PR! I just came back from PTO. I will get to this PR this week.