Refactor Affine Quantized Tensor #1234

jainapurva · 2024-11-07T01:47:52Z

Move Dtype Layout / TensorImpl class and dispatcher check-impl to their respective layout files

Affine_quantized_tensor_ops.py contains
- aten.ops implementations for aqt
- Dispatch registration and dtype kernels impl/check imports
Dtype Layout files
- torchao/dtypes/floatx
  - Float8_layout.py
  - Floatx_tensor_core_layout.py
- torchao/dtypes/uintx
  - Block_sparse_layout.py
  - Marlin_sparse_layout.py
  - Plain_layout.py
  - Semi_sparse_layout.py
  - Tensor_core_tiled_layout.py
  - Uint4_layout.py
  - uintx_layout.py
  - marlin_qqq_layout.py

Test Plan: Pass all external and internal tests

pytorch-bot · 2024-11-07T01:47:56Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/1234

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

[DomainsOnly] Jobs fail with GLIBC version not found

✅ No Failures

As of commit d992432 with merge base 06ad55a ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

torchao/dtypes/affine_quantized_tensor_ops.py

torchao/dtypes/uintx/uint8.py

torchao/dtypes/uintx/uintx.py

Ruff format and lint on some high traffic files

Update pre-commit to match CI/CD

stack-info: PR: #1228, branch: drisspg/stack/19

* add module swap UX * update * fix typing. add small notes * try NF4 support * fix * fix unpacking * fix * update nf4 integration * update backward pass

Summary: Test Plan: Reviewers: Subscribers: Tasks: Tags:

torchao/dtypes/affine_quantized_tensor_ops.py

jerryzh168 · 2024-11-14T00:52:33Z

torchao/dtypes/floatx/__init__.py

@@ -12,4 +13,6 @@
    "to_scaled_tc_floatx",
    "from_scaled_tc_floatx",
    "_SPLIT_K_MAP",
+    "Float8AQTTensorImpl",


if this is not used by other people, maybe we can remove, this is supposed to be internal mostly I think

jerryzh168 · 2024-11-14T00:53:45Z

torchao/dtypes/uintx/__init__.py

@@ -12,4 +31,11 @@
    "UintxAQTTensorImpl",
    "to_uintx",
    "_DTYPE_TO_BIT_WIDTH",
+    "_BIT_WIDTH_TO_DTYPE",
+    "UInt4Tensor",
+    "PlainAQTTensorImpl",


probably don't need to expose this one

jerryzh168 · 2024-11-14T00:54:08Z

torchao/dtypes/uintx/__init__.py

@@ -12,4 +31,11 @@
    "UintxAQTTensorImpl",


we can probably remove this one as well

jerryzh168 · 2024-11-14T00:57:47Z

torchao/quantization/quant_api.py

@@ -36,7 +36,7 @@
    to_affine_quantized_floatx_static,
    to_affine_quantized_intx,
 )
-from torchao.dtypes.uintx.uintx import UintxLayout
+from torchao.dtypes.uintx.uintx_layout import UintxLayout


we can import from torchao.dtypes for now I think

jerryzh168 · 2024-11-14T02:47:31Z

torchao/dtypes/uintx/uint4_layout.py

should be a follow up: this is not used actually, I think we can move this to prototype

Will do this in separate PR

jerryzh168 · 2024-11-14T22:25:57Z

torchao/dtypes/__init__.py

@@ -1,13 +1,9 @@
+from . import affine_quantized_tensor_ops
+
+# from ..prototype.dtypes.uint2 import UInt2Tensor, BitnetTensor


nit: we can remove this I think

jerryzh168 · 2024-11-14T22:26:04Z

torchao/dtypes/__init__.py

+from .utils import (
+    Layout,
+    PlainLayout,
+)

 # from ..prototype.dtypes.uint2 import UInt2Tensor, BitnetTensor


same for this line

jerryzh168 · 2024-11-14T22:27:35Z

test/hqq/test_hqq_affine.py

@@ -1,12 +1,7 @@
 import unittest
 import torch
-from torchao.dtypes.affine_quantized_tensor import (
-    to_affine_quantized_intx,
+from torchao.quantization.quant_primitives import (


nit: maybe we can just import from torchao.quantization now

jerryzh168 · 2024-11-14T22:29:25Z

torchao/prototype/hqq/example.py

 )
+from torchao.quantization.quant_primitives import (


jerryzh168

looks good, thanks! please import to fbcode and check if there is any internal failures

facebook-github-bot · 2024-11-15T00:57:44Z

@jainapurva has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

jainapurva requested a review from jerryzh168 November 7, 2024 01:47

facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Nov 7, 2024

jerryzh168 reviewed Nov 7, 2024

View reviewed changes

torchao/dtypes/affine_quantized_tensor_ops.py Outdated Show resolved Hide resolved

Refactored files

79a7f52

jainapurva force-pushed the aqt_refactor branch from 26e2760 to 79a7f52 Compare November 7, 2024 01:54

jerryzh168 reviewed Nov 7, 2024

View reviewed changes

torchao/dtypes/uintx/uint8.py Outdated Show resolved Hide resolved

jerryzh168 reviewed Nov 7, 2024

View reviewed changes

torchao/dtypes/uintx/uintx.py Outdated Show resolved Hide resolved

drisspg and others added 6 commits November 6, 2024 23:05

Ruff format and lint (#1226)

0cfd2b9

Ruff format and lint on some high traffic files

Update pre-commit to match CI/CD (#1227)

5d16ff1

Update pre-commit to match CI/CD

Fix for weights-only load (#1228)

43cb8d1

stack-info: PR: #1228, branch: drisspg/stack/19

Add module-swap UX for INT8 mixed-precision training (#1179)

4921bb3

* add module swap UX * update * fix typing. add small notes * try NF4 support * fix * fix unpacking * fix * update nf4 integration * update backward pass

Refactored files

235052f

Updated with ruff check

7bc89bc

jainapurva force-pushed the aqt_refactor branch from ce8934c to 7bc89bc Compare November 7, 2024 19:21

jainapurva added 4 commits November 7, 2024 11:25

<Replace this line with a title. Use 1 line only, 67 chars or less>

a060b94

Summary: Test Plan: Reviewers: Subscribers: Tasks: Tags:

Merge remote-tracking branch 'origin/main' into aqt_refactor

705a15c

Reformatted AQT

84b4d38

Merge remote-tracking branch 'origin/main' into aqt_refactor

9427064

jainapurva force-pushed the aqt_refactor branch from babdd34 to 9427064 Compare November 12, 2024 02:29

jainapurva added the topic: not user facing Use this tag if you don't want this PR to show up in release notes label Nov 12, 2024

jerryzh168 reviewed Nov 13, 2024

View reviewed changes

torchao/dtypes/affine_quantized_tensor_ops.py Outdated Show resolved Hide resolved

jainapurva added 2 commits November 13, 2024 15:39

Merge remote-tracking branch 'origin/main' into aqt_refactor

82477b9

Rebase and lint fixes

9c531e9

jainapurva force-pushed the aqt_refactor branch from 843716d to 9c531e9 Compare November 14, 2024 00:04

jainapurva marked this pull request as ready for review November 14, 2024 00:41

jerryzh168 reviewed Nov 14, 2024

View reviewed changes

jerryzh168 mentioned this pull request Nov 14, 2024

Add Int4CPULayout and update int4 woq #1278

Open

jainapurva force-pushed the aqt_refactor branch 3 times, most recently from 0315a14 to 25980e7 Compare November 14, 2024 20:32

Merge branch 'main' into aqt_refactor

af2a88d

jainapurva force-pushed the aqt_refactor branch from 25980e7 to af2a88d Compare November 14, 2024 21:59

Fixed some imports

ab8969c

jerryzh168 reviewed Nov 14, 2024

View reviewed changes

torchao/prototype/hqq/example.py Outdated

)

from torchao.quantization.quant_primitives import (

Copy link

Contributor

jerryzh168 Nov 14, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same here

jerryzh168 approved these changes Nov 14, 2024

View reviewed changes

jainapurva added 2 commits November 14, 2024 14:57

Merge remote-tracking branch 'origin/main' into aqt_refactor

6ff85f5

Minor fixes

d992432

jainapurva force-pushed the aqt_refactor branch from b9245cc to d992432 Compare November 14, 2024 23:02

jainapurva added the topic: bc-breaking Use this tag if this PR breaks backward compatibility label Nov 15, 2024

pytorch deleted a comment from facebook-github-bot Nov 15, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Refactor Affine Quantized Tensor #1234

Refactor Affine Quantized Tensor #1234

jainapurva commented Nov 7, 2024 •

edited

Loading

pytorch-bot bot commented Nov 7, 2024 •

edited

Loading

jerryzh168 Nov 14, 2024

jerryzh168 Nov 14, 2024

jerryzh168 Nov 14, 2024

jerryzh168 Nov 14, 2024

jerryzh168 Nov 14, 2024

jainapurva Nov 14, 2024

jerryzh168 Nov 14, 2024

jerryzh168 Nov 14, 2024

jerryzh168 Nov 14, 2024

jerryzh168 Nov 14, 2024

jerryzh168 left a comment

facebook-github-bot commented Nov 15, 2024

		@@ -1,13 +1,9 @@
		from . import affine_quantized_tensor_ops

		# from ..prototype.dtypes.uint2 import UInt2Tensor, BitnetTensor

Refactor Affine Quantized Tensor #1234

Are you sure you want to change the base?

Refactor Affine Quantized Tensor #1234

Conversation

jainapurva commented Nov 7, 2024 • edited Loading

pytorch-bot bot commented Nov 7, 2024 • edited Loading

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/1234

❗ 1 Active SEVs

✅ No Failures

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jerryzh168 left a comment

Choose a reason for hiding this comment

facebook-github-bot commented Nov 15, 2024

jainapurva commented Nov 7, 2024 •

edited

Loading

pytorch-bot bot commented Nov 7, 2024 •

edited

Loading