Add TorchAOHfQuantizer
Author: jerryzh168
Pull request: yes
Labels: name: Quantization

This PR adds [torchao](https://github.com/pytorch/ao) as an option for quantization, right now we expose three quantization technique: `int4_weight_only`, `int8_weight_only` and `int8_dynamic_activation_int8_weight` to user, it uses tensor subclass to implement weight quantization and `torch.compile` to generate efficient code for cuda.


Example:
```
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers import TorchAoConfig
from torchao.utils import benchmark_model
import torchao

DEVICE_TYPE = "cuda"

def init_model_and_benchmark(model_id, torch_dtype=torch.bfloat16, quantization_config=None):
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    if quantization_config is not None:
        model = AutoModelForCausalLM.from_pretrained(model_id, device_map=DEVICE_TYPE, torch_dtype=torch.bfloat16, quantization_config=quantization_config)
    else:
        model = AutoModelForCausalLM.from_pretrained(model_id, device_map=DEVICE_TYPE, torch_dtype=torch.bfloat16)

    # sanity check: run the model
    input_text = "What are we having for dinner?"
    input_ids = tokenizer(input_text, return_tensors="pt").to(DEVICE_TYPE)
    output = model.generate(**input_ids, max_new_tokens=1000)
    print(tokenizer.decode(output[0], skip_special_tokens=True))

    NUM_WARMUP = 10
    NUM_RUNS = 50

    if quantization_config is not None:
        torchao.quantization.utils.recommended_inductor_config_setter()

    model = torch.compile(model, mode="max-autotune")

    benchmark_model(model.generate, NUM_WARMUP, kwargs=input_ids, device_type=DEVICE_TYPE)
    print("running benchmark")
    results = benchmark_model(model.generate, NUM_RUNS, kwargs=input_ids, device_type=DEVICE_TYPE)
    return model, results

model_id = "meta-llama/Llama-2-7b-chat-hf"
bf16_model, bf16_time = init_model_and_benchmark(model_id)
print(f"bf16: {bf16_time}")

quantization_config = TorchAoConfig(quant_type="int4_weight_only", group_size=32)
int4wo_model, int4wo_time = init_model_and_benchmark(model_id)
print(f"int4wo: {int4wo_time} speedup: {bf16_time / int4wo_time}")

quantization_config = TorchAoConfig(quant_type="int8_weight_only")
int8wo_model, int8wo_time = init_model_and_benchmark(model_id)
print(f"int8wo: {int8wo_time} speedup: {bf16_time / int8wo_time}")

quantization_config = TorchAoConfig(quant_type="int8_dynamic_activation_int8_weight")
int8da_int8w_model, int8da_int8w_time = init_model_and_benchmark(model_id)
print(f"int8da_int8w: {int8da_int8w_time} speedup: {bf16_time / int8da_int8w_time}")
```

 -->
Perf benchmark results (WARMUP=2, RUNS=10) A100 cuda machine:

```
bf16: 8167.34453125
int4wo: 6036.775390625 speedup: 1.3529316568467553
int8wo: 6974.3453125 speedup: 1.1710553701165052
int8da_int8w: 7945.5375 speedup: 1.0279159252913979
```

we don't always expect speedup here since we may need manual changes to model code before we can get speedup

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_32306). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

cc @SunMarc 

> Hi @jerryzh168, thanks for doing the integration ! I did a first pass and the pr looks pretty clean ! I've left a few comments. Also, can you share a bit about how serialization might look like and the issue that you are facing ? It's fine if we don't support it in this PR. This can come in a follow-up PR !

@SunMarc thanks for the review, I have addressed/replied to all the comments.

I talked about serialization briefly in https://github.com/jerryzh168/transformers/blob/torchao/docs/source/en/quantization/torchao.md

main thing is torchao is using tensor subclasses, and it does not work well with serialization code in huggingface, we have opened issue here: https://github.com/huggingface/transformers/issues/32364, the problem we are seeing now is huggingface is relying on `storage_id` to figure out the sharing relationship of tensors: https://github.com/huggingface/huggingface_hub/blob/97d5ef603f41314a52eb2d045ec966cf9fed6295/src/huggingface_hub/serialization/_torch.py#L479, and the assumption that all tensors have storage is been broken by subclassed tensors. we can recommend a more general solutions for this. cc @albanD



@SunMarc thanks for the review, I addressed your comments, and had one remaining question (see inline)

Is safetensors serialisation supported for the quantized weights?

Reference: https://github.com/huggingface/safetensors/issues/515#issuecomment-2271208072

> Is safetensors serialisation supported for the quantized weights?
> 
> Reference: [huggingface/safetensors#515 (comment)](https://github.com/huggingface/safetensors/issues/515#issuecomment-2271208072)

safetensor is not supported yet, still figuring out the plan there, but I'd expect us to support non-safe tensor serialization first

Hey @SunMarc anything else we need to get this landed? Sort of related we're releasing ao 0.4 tomorrow and was hoping we could add this integration to our release notes

Here's the PR with a couple of fix: https://github.com/jerryzh168/transformers/pull/1. Oh my bad, I pushed on your branch by mistake and it merged the PR. I didn't know I could push ;) Feel free to revert if you don't agree with some changes I made ! 

> Here's the PR with a couple of fix: [jerryzh168#1](https://github.com/jerryzh168/transformers/pull/1). Oh my bad, I pushed on your branch by mistake and it merged the PR. I didn't know I could push ;) Feel free to revert if you don't agree with some changes I made !

I see, that's fine, thanks for the fix! I thought `get_keys_to_not_convert` is model specific? does it apply to all models?

`get_keys_to_not_convert` applies to all models. This is something we added for all our quantization methods and it seems to work pretty well ! I also changed a bit the torch_dtype logic ! 

I also fixed a new nits in the tests !

Hi everyone, I ran some experiments and it is actually possible to train a torchao-quantized model with PEFT, even though it is not explicitly supported in PEFT yet (working on it). The reason why it works is because the type of the layer is not changed by torchao, so for instance `nn.Linear` stays `nn.Linear`. Therefore, PEFT will apply a normal PEFT layer, e.g. `lora.Linear`.

I wrote some unit tests to check this and it looks like it works:

```python
import gc
import tempfile
import unittest

import torch
from datasets import load_dataset
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from transformers import AutoModelForCausalLM, AutoTokenizer, DataCollatorForLanguageModeling, TrainingArguments, Trainer

class PeftTorchaoGPUTests(unittest.TestCase):
    supported_quant_types = [
        "int8_weight_only",
        "int8_dynamic_activation_int8_weight",
    ]

    def setUp(self):
        self.causal_lm_model_id = "facebook/opt-125m"
        self.tokenizer = AutoTokenizer.from_pretrained(self.causal_lm_model_id)

    def tearDown(self):
        gc.collect()
        torch.cuda.empty_cache()

    def test_causal_lm_training_single_gpu_torchao(self):
        # call this with CUDA_VISIBLE_DEVICES=0
        from transformers import TorchAoConfig

        device = 0

        for quant_type in self.supported_quant_types:
            with tempfile.TemporaryDirectory() as tmp_dir:
                quantization_config = TorchAoConfig(quant_type=quant_type)
                model = AutoModelForCausalLM.from_pretrained(
                    self.causal_lm_model_id, device_map=device, quantization_config=quantization_config
                )
                model = prepare_model_for_kbit_training(model)

                config = LoraConfig(
                    r=16,
                    lora_alpha=32,
                    target_modules=["q_proj", "v_proj"],
                    lora_dropout=0.05,
                    bias="none",
                    task_type="CAUSAL_LM",
                )
                model = get_peft_model(model, config)

                data = load_dataset("ybelkada/english_quotes_copy")
                data = data.map(lambda samples: self.tokenizer(samples["quote"]), batched=True)

                trainer = Trainer(
                    model=model,
                    train_dataset=data["train"],
                    args=TrainingArguments(
                        per_device_train_batch_size=4,
                        gradient_accumulation_steps=4,
                        warmup_steps=2,
                        max_steps=5,
                        learning_rate=2e-4,
                        logging_steps=1,
                        output_dir=tmp_dir,
                    ),
                    data_collator=DataCollatorForLanguageModeling(self.tokenizer, mlm=False),
                )
                model.config.use_cache = False
                trainer.train()

    def test_causal_lm_training_multi_gpu_torchao(self):
        from transformers import TorchAoConfig

        for quant_type in self.supported_quant_types:
            with tempfile.TemporaryDirectory() as tmp_dir:
                quantization_config = TorchAoConfig(quant_type=quant_type)
                model = AutoModelForCausalLM.from_pretrained(
                    self.causal_lm_model_id,
                    device_map="auto",
                    quantization_config=quantization_config,
                    torch_dtype=torch.bfloat16,
                )

                assert set(model.hf_device_map.values()) == set(range(torch.cuda.device_count()))

                model = prepare_model_for_kbit_training(model)
                setattr(model, "model_parallel", True)
                setattr(model, "is_parallelizable", True)

                config = LoraConfig(
                    r=16,
                    lora_alpha=32,
                    target_modules=["q_proj", "v_proj"],
                    lora_dropout=0.05,
                    bias="none",
                    task_type="CAUSAL_LM",
                )
                model = get_peft_model(model, config)

                data = load_dataset("Abirate/english_quotes")
                data = data.map(lambda samples: self.tokenizer(samples["quote"]), batched=True)

                trainer = Trainer(
                    model=model,
                    train_dataset=data["train"],
                    args=TrainingArguments(
                        per_device_train_batch_size=4,
                        gradient_accumulation_steps=4,
                        warmup_steps=2,
                        max_steps=5,
                        learning_rate=2e-4,
                        logging_steps=1,
                        output_dir=tmp_dir,
                    ),
                    data_collator=DataCollatorForLanguageModeling(self.tokenizer, mlm=False),
                )
                model.config.use_cache = False
                trainer.train()
```

As you can see, `int8_weight_only` and `int8_dynamic_activation_int8_weight` work already. For `int4_weight_only`, I got an error unfortunately:

```
tests/test_gpu_examples.py:3079: in test_causal_lm_training_single_gpu_torchao
    trainer.train()
../../clones/transformers/src/transformers/trainer.py:1991: in train
    return inner_training_loop(
../../clones/transformers/src/transformers/trainer.py:2327: in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs)
../../clones/transformers/src/transformers/trainer.py:3452: in training_step
    self.accelerator.backward(loss, **kwargs)
../../../anaconda3/envs/peft/lib/python3.11/site-packages/accelerate/accelerator.py:2159: in backward
    loss.backward(**kwargs)
../../../anaconda3/envs/peft/lib/python3.11/site-packages/torch/_tensor.py:521: in backward
    torch.autograd.backward(
../../../anaconda3/envs/peft/lib/python3.11/site-packages/torch/autograd/__init__.py:289: in backward
    _engine_run_backward(
../../../anaconda3/envs/peft/lib/python3.11/site-packages/torch/autograd/graph.py:768: in _engine_run_backward
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
../../../anaconda3/envs/peft/lib/python3.11/site-packages/torch/autograd/function.py:306: in apply
    return user_fn(self, *args)
../../../anaconda3/envs/peft/lib/python3.11/site-packages/torch/utils/checkpoint.py:313: in backward
    torch.autograd.backward(outputs_with_grad, args_with_grad)
../../../anaconda3/envs/peft/lib/python3.11/site-packages/torch/autograd/__init__.py:289: in backward
    _engine_run_backward(
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

t_outputs = (tensor([[[-0.4355, -0.8750,  0.2188,  ..., -0.8750, -0.4375,  0.0693],
         [ 0.0540, -1.1094,  0.0889,  ...,  0....,  0.2139,  ..., -0.5117, -0.3672,  0.2539]]],
       device='cuda:0', dtype=torch.bfloat16, grad_fn=<ViewBackward0>),)
args = ((tensor([[[-8.0466e-06, -4.8876e-05, -2.2888e-05,  ...,  6.0320e-05,
          -1.9073e-06, -9.4414e-05],
         [-...,  ...,  0.0000e+00,
           0.0000e+00,  0.0000e+00]]], device='cuda:0', dtype=torch.bfloat16),), False, False, ())
kwargs = {'accumulate_grad': True, 'allow_unreachable': True}, attach_logging_hooks = False

    def _engine_run_backward(t_outputs, *args, **kwargs):
        attach_logging_hooks = log.getEffectiveLevel() <= logging.DEBUG
        if attach_logging_hooks:
            unregister_hooks = _register_logging_hooks_on_whole_graph(t_outputs)
        try:
>           return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
                t_outputs, *args, **kwargs
            )  # Calls into the C++ engine to run the backward pass
E           RuntimeError: derivative for aten::_weight_int4pack_mm is not implemented

```

For me to be able to use `Trainer`, I had to patch `is_trainable` to return `True`. So I wonder if it can be set to `True` if PEFT is available and int8 is being used. WDYT?

Note that I'll still look into a proper PEFT integration, since I'm not sure if features like merging would already work. But training and inference should be fine as is.

PS: I'm using torch 2.4.0, PEFT and transformers from source, in case it matters.

Nice ! I'll create to PR for that 

Hello. The current doc has strange content about the features that torchao supports. It would be good if someone can fix it.

https://huggingface.co/docs/transformers/quantization/overview

![图片](https://github.com/user-attachments/assets/49a49181-ffc9-4ffa-ba9e-cdee24ef595f)


@imba-tjd which one are you referring to?

@jerryzh168 Please see the screenshot in my above comment.
For example there is `🟢🔴` in one cell at the line of torchao.

> @jerryzh168 Please see the screenshot in my above comment. For example there is `🟢🔴` in one cell at the line of torchao.

yeah this means partial support. but it might be clearer is we make it yellow and then add a note I guess, specifically, this means we support non-safetensor serialization (please feel free to submit a fix if you have time, otherwise I should be able to get to this a bit later)