Add TorchAOHfQuantizer

[jerryzh168]

This PR adds torchao 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

[msaroufim]

We do have support for nf4 and QAT though, is the intent here to say the hf ao backend does not support this yet?

[jerryzh168]

QAT is still in prototype I think, oh I forgot about nf4, I haven't thought of how we expose nf4 yet, will check with Driss

[jerryzh168]

I think we haven't tested nf4 dtype with quantize_ API yet, maybe we can create an issue to make nf4 tensor work with quantize_ API

[jerryzh168]

checked with Driss, we probably need to do a separate integration for nf4/QLoRA like what we did for torchtune, need to implement a (Q)LoRALinear class like this https://github.com/pytorch/torchtune/blob/288ff4435b0cf17325b5c3b112f6859a6cdf0ea2/torchtune/modules/peft/lora.py#L18 in huggingface

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[amyeroberts]

cc @SunMarc

[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]

I think that we can the user about that in update_torch_dtype

[jerryzh168]

you mean update the default dtype to torch.bfloat16?

[jerryzh168]

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: #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 iterating @jerryzh168 ! We are practically there ! Left a few comments

[SunMarc]

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 !

[SunMarc]

Thanks for iterating @jerryzh168 ! LGTM ! Excited to have torchao quantization methods in transformers !

[SunMarc]

I also fixed a new nits in the tests !

[sayakpaul]

But one can serialize the state dict with quantized_model.state_dict() and then use it later right? I'd prefer to have this information included as a workaround.

Also, can we provide a small table noting the expected memory and latency savings here?

[jerryzh168]

yeah state_dict works locally, but I'm not sure how to get this work with the load_pretrained save_pretrained API. we are adding serialization support for both safetensor and non-safetensor though

yeah we have some numbers here: https://github.com/pytorch/ao/tree/main/torchao/quantization#benchmarks, I can add it

[sayakpaul]

Are we not supporting fp8?

[msaroufim]

fp8 inference is a bit more experimental, we can add it in a later PR

[ArthurZucker]

Thanks for this nice work both! 🤗

[ArthurZucker]

Let's maybe mention it supports compile ?

[jerryzh168]

added above

[ArthurZucker]

I am guessing that the type of the layer under torchao does not change the linear layer? As in any other quantization methods we check against the quantized liner type!

[SunMarc]

That's right, torchao only modifies the weight tensor and not the nn.Linear module. So we don't need to check against the quantized linear type.

[BenjaminBossan]

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:

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.

[SunMarc]

Nice ! I'll create to PR for that