Faster generation using AWQ + Fused modules

[younesbelkada]

What does this PR do?

Introduces a new feature - fused module generation using autoawq library. Users need to specify modules that they want to fuse inside fusing_mapping.

The API is as follows:

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, AwqConfig, TextStreamer

model_id = "TheBloke/Mistral-7B-OpenOrca-AWQ"
torch_device = "cuda" if torch.cuda.is_available() else "cpu"

quantization_config = AwqConfig(
    bits=4,
    do_fuse=True,
    fuse_max_seq_len=512,
)

model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=quantization_config).to(0)
tokenizer = AutoTokenizer.from_pretrained(model_id)

streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

prompt_template = """\
<|im_start|>system
You are MistralOrca, a large language model trained by Alignment Lab AI. Write out your reasoning step-by-step to be sure you get the right answers!<|im_end|>
<|im_start|>user
{prompt}<|im_end|>
<|im_start|>assistant"""

prompt = "You're standing on the surface of the Earth. "\
        "You walk one mile south, one mile west and one mile north. "\
        "You end up exactly where you started. Where are you?"

tokenizer.pad_token = tokenizer.eos_token

inputs = tokenizer([prompt_template.format(prompt=prompt), prompt_template.format(prompt=prompt), prompt_template.format(prompt=prompt)], return_tensors="pt", padding=True).to(0)

outputs = model.generate(**inputs, max_new_tokens=512)
print(tokenizer.batch_decode(outputs, skip_special_tokens=True))

Before this PR:

After this PR:

TODOs:

•  Block save_pretrained if one creates a fused model
• Fuse MLP
• wait until dynamic batch size is resolved on auto-awq side : Adaptive batch sizing casper-hansen/AutoAWQ#181
• Force users to use autoawq>=0.1.7
• Tests
• Benchmarks
• Docs

cc @amyeroberts @casper-hansen @SunMarc

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.

[younesbelkada]

Before moving forward with tests and advanced docs, I would love to have an early feedback of the API that is described in the PR description. cc @amyeroberts , whenever you have time, i would appreciate your feedback on this PR 🙏 Thanks!

[amyeroberts]

Really nice piece of work! 🔥

Main comment is about structure of the input arguments. It might just be my current understanding from the PR so please do correct me if I'm wrong.

At the moment, it seems there's two types of arguments which configure the behaviour: the fusing_mapping and max_seq_len. Does max_seq_len have its own config value because we expect this to be the value users modify a lot?

For AwqConfig do we foresee other fuse arguments to be added to configure this behaviour? If so, we might want to bundle them all together into a single fuse_config which the AwqConfig owns.

[amyeroberts]

If hidden_size, num_attentions_heads and num_key_value_heads are required arguments it would be good to verify these keys exist in the mapping from quantization_config.fusion_mapping

[younesbelkada]

I propose to do it in 0a08551 lmk what do you think!

[amyeroberts]

Thanks!

[amyeroberts]

This line here indicates a possible area of brittleness in the current design: there's going to be many different variations of model attribute names which - if this becomes popular - we'll have to account for. Passing in quantization_config.fusing_mapping is great to provide flexibility for the mapping.

My understanding is that quantization_config.fuse_max_seq_len is passed in separately because the value will be specific to each model's configuration rather than e.g. architecture. The question is - why are params like "hidden_size" and "num_key_value_heads" not passed in with quantization_config.fusing_mapping? I think it would make more sense for them all to be passed in together or have two separate dictionaries - one to map layer names and the other to map model-specific configs.

[younesbelkada]

I proposed to do that to handle different architecture variants , 7b, 13b, .. etc. I did not find a way to properly map that in AWQ_FUSED_MAPPINGS and decided to make it as architecture agnostic as possible.
For now we only support mistral and llama if one passes fuse_modules, for other architectures users need to manually create a mapping and pass it through AwqConfig. If we want to support other architectures in the future, and face attribute errors I propose to fix it directly in the corresponding config object by adding an attribute_map: https://github.com/huggingface/transformers/blob/main/src/transformers/models/t5/configuration_t5.py#L83 to make sure num_attention_heads, and hidden_size exists. What do you think?

[amyeroberts]

Sounds good to me!

It handles almost all cases and if there's anything more complex in the future we can handle it then when we know more about the problem.

[amyeroberts]

This is a pretty big function - I'd split it up so that there's private functions for each layer-replacement which is then called within the big for-loop

[younesbelkada]

Done in cde53ef

[amyeroberts]

Beautiful 🤩

[amyeroberts]

Can we guarantee fusing_mapping has an "mlp" key? and that it has at least one value?

[younesbelkada]

I propose to do a more extensive check in b187c07 and I changed the logic there a bit to make sure these methods will do nothing in case we put an empty array in these fields

[amyeroberts]

gate_, up_ down_ is highly specific to just a few models. Instead we can generalise this to take a list of linear layers so that any MLP can be passed

[amyeroberts]

Same comment here re config param names. hidden_act is common but it's not always the name used. Could we pass this in in the fuse config too?

[younesbelkada]

Hmm, do you think my suggestion here: #27411 (comment) could be applied in this case?

[amyeroberts]

Thanks for iterating!

Great code, great tests, great docs 💯

[amyeroberts]

Sounds good to me!

It handles almost all cases and if there's anything more complex in the future we can handle it then when we know more about the problem.

[amyeroberts]

Thanks!

[amyeroberts]

Beautiful 🤩

[amyeroberts]

Based on these keys it would be cool to have a tool which automatically generates a config for a model assuming you wanted to fuse all modules.

[SunMarc]

Thanks for making AWQ faster through fused modules 🔥. The design looks great and could be easily extended to other quantization scheme in the future. I left a few comments.

[younesbelkada]

Thanks @amyeroberts @SunMarc for your great reviews!
@SunMarc I just want to get more clarification on this comment, otherwise good to merge IMO !

[SunMarc]

Thanks for iterating on this ! I've left a few comments about two points but they are not blocking for this PR !