| 1 | # TorchAO |
| 2 | |
| 3 | TorchAO is an architecture optimization library for PyTorch, it provides high performance dtypes, optimization techniques and kernels for inference and training, featuring composability with native PyTorch features like torch.compile, FSDP etc.. Some benchmark numbers can be found [here](https://github.com/pytorch/ao/tree/main/torchao/quantization#benchmarks). |
| 4 | |
| 5 | We recommend installing the latest torchao nightly with |
| 6 | |
| 7 | ```bash |
| 8 | # Install the latest TorchAO nightly build |
| 9 | # Choose the CUDA version that matches your system (cu126, cu128, etc.) |
| 10 | pip install \ |
| 11 | --pre torchao>=10.0.0 \ |
| 12 | --index-url https://download.pytorch.org/whl/nightly/cu126 |
| 13 | ``` |
| 14 | |
| 15 | ## Quantizing HuggingFace Models |
| 16 | |
| 17 | You can quantize your own huggingface model with torchao, e.g. [transformers](https://huggingface.co/docs/transformers/main/en/quantization/torchao) and [diffusers](https://huggingface.co/docs/diffusers/en/quantization/torchao), and save the checkpoint to huggingface hub like [this](https://huggingface.co/jerryzh168/llama3-8b-int8wo) with the following example code: |
| 18 | |
| 19 | ??? code |
| 20 | |
| 21 | ```Python |
| 22 | import torch |
| 23 | from transformers import TorchAoConfig, AutoModelForCausalLM, AutoTokenizer |
| 24 | from torchao.quantization import Int8WeightOnlyConfig |
| 25 | |
| 26 | model_name = "meta-llama/Meta-Llama-3-8B" |
| 27 | quantization_config = TorchAoConfig(Int8WeightOnlyConfig()) |
| 28 | quantized_model = AutoModelForCausalLM.from_pretrained( |
| 29 | model_name, |
| 30 | dtype="auto", |
| 31 | device_map="auto", |
| 32 | quantization_config=quantization_config |
| 33 | ) |
| 34 | tokenizer = AutoTokenizer.from_pretrained(model_name) |
| 35 | input_text = "What are we having for dinner?" |
| 36 | input_ids = tokenizer(input_text, return_tensors="pt").to("cuda") |
| 37 | |
| 38 | hub_repo = # YOUR HUB REPO ID |
| 39 | tokenizer.push_to_hub(hub_repo) |
| 40 | quantized_model.push_to_hub(hub_repo, safe_serialization=False) |
| 41 | ``` |
| 42 | |
| 43 | Alternatively, you can use the [TorchAO Quantization space](https://huggingface.co/spaces/medmekk/TorchAO_Quantization) for quantizing models with a simple UI. |