Add torchao to optimum as a pytorch backend configuration

[jerryzh168]

Summary:
att.
for now we just added int4 weight only quantization, using TorchAoConfig from https://huggingface.co/docs/transformers/main/en/quantization/torchao

Test Plan:
python examples/pytorch_llama.py

gpt2 model

Base:
prefill: 9486.25 decode: 83.75

AWQ:
prefill: 9496.02 decode: 83.62

GPTQ
prefill: 9814.27 decode: 97.31

torchao int4wo
prefill: 10007.93 decode: 84.66

llama2

Base:
prefill: 2275.32 decode: 18.92

AWQ:
prefill: 2344.19 decode: 18.21

GPTQ:
prefill: 2881.87 decode: 26.47

torchao int4wo
prefill: 3035.82 decode: 24.51

Reviewers:

Subscribers:

Tasks:

Tags:

[IlyasMoutawwakil]

@jerryzh168 great prefill performance !
can you also add torchao to the installed libs in our docker images.

[jerryzh168]

@IlyasMoutawwakil I feel we should be able to get better performance if qkv are fused in llama model, which is what people typically do today, not sure if there is a huggingface model in the hub that has it though

by adding torchao to docker, do you mean these: https://github.com/huggingface/optimum-benchmark/tree/main/docker ?

[IlyasMoutawwakil]

@IlyasMoutawwakil I feel we should be able to get better performance if qkv are fused in llama model, which is what people typically do today, not sure if there is a huggingface model in the hub that has it though

Would require changes in transformers modeling to support that, would make for a great PR there, some frameworks (like ipex in optimum-intel) do that manually.

by adding torchao to docker, do you mean these: https://github.com/huggingface/optimum-benchmark/tree/main/docker ?

yes exactly, the same as torchvision and torchaudio, I would rather have torchao installed there to avoid any cpu/cuda version issues.

[jerryzh168]

yeah we are doing that manually in our benchmarks as well, maybe we can worry about it a bit later. I'll make docker changes first

[jerryzh168]

@IlyasMoutawwakil I added torchao to cpu and cuda dockers, it's not available for rocm 5.7, but it's available for 6.1 and 6.2: https://download.pytorch.org/whl/nightly/torchao/