Migrating from AffineQuantizedTensor + Layouts to new structure of tensor subclasses #2752
Description
Update: Our team will evaluate this more before outsourcing the migration to more people in the community
Context:
Previously we use AffineQuantizedTensor for many of our use cases including int4, float8, intx, floatx. It introduces some complicated abstractions like Layout, people have been saying it's a bit hard to understand, and there are many indirections in the code.
As an effort simplify the code base and make it easier to contribute to, we have been adding new features with a different structure in mind. Now we want to structure Tensors by "dtype" and "packing_format", e.g. we'll have Int4PreshuffledTensor, Int8Tensor, Float8Tensor, instead of having AffineQuantizedTensor and multiple layouts.
Please check out our updated docs for the new tensor subclass organization structure and guide for design:
- quantization overview: https://docs-preview.pytorch.org/pytorch/ao/2723/quantization_overview.html
- contributor guide: https://docs-preview.pytorch.org/pytorch/ao/2723/contributor_guide.html
- Examples of tensor subclasses following new design: https://github.com/pytorch/ao/tree/main/torchao/quantization/quantize_/workflows
migration status
| inference config name | current status | plan | POC | status |
|---|---|---|---|---|
MXFPInferenceConfig |
built on v2 | n/a | - | done |
NVFP4InferenceConfig |
built on v2 | n/a | - | done |
Float8DynamicActivationInt4WeightConfig |
built on v2 | n/a | - | done |
Int4WeightOnlyConfig |
v2 and v1 exists | deprecate v1 | ? | #3513 |
Int8DynamicActivationIntxWeightConfig |
v2 and v1 exists | deprecate v1 | ? | #3511 |
Float8WeightOnlyConfig |
v2 and v1 exists | deprecate v1 | ? | #3510 |
Float8DynamicActivationFloat8WeightConfig |
v2 and v1 exists | deprecate v1 | ? | #3510 |
IntxWeightOnlyConfig |
v2 and v1 exists | deprecate v1 | ? | #3512 |
Float8DynamicActivationFloat8SemiSparseWeightConfig |
v1 exists | add a new sparse packing_format for float8 dynamic quant config, then deprecate v1 | ? | #3361 |
Int8WeightOnlyConfig |
v1 exists | create v2, then deprecate v1 | ? | #3407 |
Int8DynamicActivationInt8WeightConfig |
v1 exists | create v2, then deprecate v1 | ? | #3407 |
Int8DynamicActivationInt4WeightConfig |
v1 exists | move to prototype | ? | #3491 |
Int4DynamicActivationInt4WeightConfig |
v1 exists | move to prototype | ? | #3491 |
GemliteUIntXWeightOnlyConfig |
v1 exists | move to prototype | ? | #3491 |
Float8StaticActivationFloat8WeightConfig |
v1 exists | move to prototype | ? | #3491 |
UIntXWeightOnlyConfig |
v1 exists | move to prototype | ? | #3491 |
FPXWeightOnlyConfig |
v1 exists | move to prototype | ? | #3491 |
appendix
List of things to migrate:
INT8
- [move to prototype] https://github.com/pytorch/ao/blob/main/torchao/dtypes/uintx/block_sparse_layout.py @jainapurva Move block_sparse_layout to prototype #3276
- [migrate] https://github.com/pytorch/ao/blob/main/torchao/dtypes/uintx/plain_layout.py @namgyu-youn introduce new int8 quantization API #3241
- [move to prototype] https://github.com/pytorch/ao/blob/main/torchao/dtypes/uintx/semi_sparse_layout.py @namgyu-youn Introduce new W8A8-FP-CSR quantitzation API #3258 (no need to migrate to new tensor structure)
[migration done, TODO: delete old path after all migration is done] INT4 weight only
- https://github.com/pytorch/ao/blob/main/torchao/dtypes/uintx/int4_cpu_layout.py @Xia-Weiwen https://github.com/pytorch/ao/blob/main/torchao/quantization/quantize_/workflows/int4/int4_opaque_tensor.py
- https://github.com/pytorch/ao/blob/main/torchao/dtypes/uintx/int4_xpu_layout.py @liangan1 Add Int4PlainInt32Tensor #2845
- https://github.com/pytorch/ao/blob/main/torchao/dtypes/uintx/marlin_sparse_layout.py @liangel-02 Int4 sparse marlin tensor #2771
- https://github.com/pytorch/ao/blob/main/torchao/dtypes/uintx/tensor_core_tiled_layout.py @jerryzh168 Add Int4TilePackedTo4dTensor #2791
- HQQ support for tensor core tiled layout @jerryzh168 Add hqq support for Int4TilePackedTo4dTensor #2912
[move to prototype] INT4 weight + int8 activation
- https://github.com/pytorch/ao/blob/main/torchao/dtypes/uintx/cutlass_int4_packed_layout.py @jainapurva Move cutlass_int4_packed_layout to prototype #3277
- https://github.com/pytorch/ao/blob/main/torchao/dtypes/uintx/dyn_int8_act_int4_wei_cpu_layout.py @jainapurva Move dyn_int8_act_int4_wei_cpu_layout to prototype/dtypes #3299
- https://github.com/pytorch/ao/blob/main/torchao/dtypes/uintx/marlin_qqq_tensor.py @jainapurva Move marlin_qqq_tensor to prototype/dtypes #3307
UINTx Weight Only
- [move to protoype] https://github.com/pytorch/ao/blob/main/torchao/dtypes/uintx/gemlite_layout.py @jainapurva Move gemlite layout to prototype/dtypes #3313
- [move to protoype] https://github.com/pytorch/ao/blob/main/torchao/dtypes/uintx/uintx_layout.py @jainapurva Move uintx_layout to prototype/dtypes #3316
[migration done, TODO: delete old path after all migration is done] Int8DynamicActivationIntxWeightConfig
- https://github.com/pytorch/ao/blob/main/torchao/dtypes/uintx/packed_linear_int8_dynamic_activation_intx_weight_layout.py @metascroy Introduce IntxOpaqueTensor to replace PackedInt8DynamicActivationIntxWeightLayout in AQT #2742
- https://github.com/pytorch/ao/blob/main/torchao/dtypes/uintx/q_dq_layout.py @metascroy Add IntxUnpackedTensor #2732
FP8
- [migrate] https://github.com/pytorch/ao/blob/main/torchao/dtypes/floatx/cutlass_semi_sparse_layout.py @namgyu-youn Introduce new W8A8-FP-CSR quantitzation API #3258 and @bbeckca [WIP] Move float8 cutlass sparse layout to Float8SemiSparseTensor #3182
FPx