[Tracker] WIP features for torchao 0.3

[supriyar]

Focus - benchmarking, documentation, tutorials, prototype to beta

Due date: June 13 2024

Spillover from 0.2.0

• Consolidating workflows to use tensor subclass @jerryzh168
• Fast sparse training @jcaip
• Explore adding HQQ 4/3/2-bit quant to torchao @HDCharles @mobicham
• Don't cause a "import error" when someone is using a feature unsupported by e.g. torch 2.1.2 (e.g. https://github.com/pytorch-labs/ao/blob/046dc985de6d5eac05c6575cc71505687e3aadf1/torchao/quantization/quant_primitives.py#L42 will cause an import error if someone tries to use torchao.quantization.quant_primitives.per_token_dynamic_quant on 2.2.2

Benchmarking

• Setup model level benchmark for accuracy and performance in torchbench for single quantization API, so that we can start deprecating quant primitives and quant apis after making sure no regressions (@HDCharles)
• Benchmarks for auto quant on pytorch benchmark's inference quant pane (@HDCharles)

Documentation

• Docs Revamp #181
• Make our website sparkle https://pytorch.org/ao (everyone)
• Unify on README or website, but don't duplicate content
• Get started on a design to slot in the prototype features - HQQ (@msaroufim)

Tutorials

• Tutorial for affine quantization dtype and unified quant primitives - Found lots of subtle differences, especially w.r.t. preserving zeros and tinygemm (@jerryzh168)

Core

• QAT workflow (@andrewor14)
• dedup the implementations of quant primitives (@jerryzh168)
• dedup the implementations of quant APIs (@jerryzh168)
• Deduplicate int4 workflows
• Factory function ahd implements decorator for affine quantization dtype
• Bit packing interfaces @msaroufim
• float6 kernels @gau-nernst
• int 3/5 kernel @msaroufim

[jeromeku]

@msaroufim

• Generally, what needs to be done to compose a new dtype with FSDP?
• What other (high priority) dtypes are on the ao roadmap for integration with FSDP?
• Is there a universal representation for asymmetric / symmetrically quantized types in torch? I.e., a subbyte/byte type with scale / zero that can be used regardless of the quantization method?
• Is development of fp8 primitives for training and inference primarily in pytorch/float8_experimental or are there specific torchao initiatives focused on fp8?

Happy to contribute on any of these fronts.

[jerryzh168]

@jeromeku for

• Is there a universal representation for asymmetric / symmetrically quantized types in torch? I.e., a subbyte/byte type with scale / zero that can be used regardless of the quantization method?

yes, it's called AffineQuantizedTensor, we are putting all types of symmetric/asymmetric, per_tensor/channel/group/token, int8/int4/int3/... under this tensor subclass.

here is an model level API walk through using the tensor subclass: https://github.com/pytorch/ao/tree/main/torchao/quantization#quantization-flow

currently I'm working on replacing existing APIs with it: #294, after that I also plan to publish a more detailed tutorials to talk about how to implement a new data representation with tensor subclass, using this as an example

[msaroufim]

@jeromeku regarding your other questions

• To compose new dtypes with FSDP you can follow the playbook here [FSDP2][NF4Tensor][2/n] implement torch.chunk and other ops #150 we'll look to write some docs @weifengpy but let us know if you have any questions in the meantime
• I guess regarding high pri dtypes, I'm not 100% sure yet cause it depends on what researchers might do - I'm personally biased to getting bitnet to work
• Regarding fp8 training we're looking to centralize the fp8 work in this repo and @vkuzo is gonna be moving bits over time

[bhack]

Are we going to support dynamic inputs?

[msaroufim]

Hi @bhack! Seen you on a lot of threads. Could you share a bit more what you mean by dynamic inputs? Are you referring to dynamic shapes?

[bhack]

Yes dynamic input shape but not mainly on the batch dimension. So e.g. image with different width and height