-
Notifications
You must be signed in to change notification settings - Fork 260
Add Float8QuantizedTensor (AQT subclass) and replace to_affine_quantized_floatx with to_affine_quantized_float8 in quantization APIs #1599
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/1599
Note: Links to docs will display an error until the docs builds have been completed. ❌ 2 New FailuresAs of commit 2f15cc1 with merge base 32d9b0b ( NEW FAILURES - The following jobs have failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
thanks, we also want to split out a Float8 (and floatx) specific AQT implementations as well, I talked to @jainapurva before |
Yep that makes sense, when I talked to her earlier she said she is planning to create these AQT subclasses, so I decided to do this part of the refactor. |
torchao/dtypes/__init__.py
Outdated
@@ -38,6 +34,7 @@ | |||
"to_affine_quantized_fpx", | |||
"to_affine_quantized_floatx", |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Please remove floatx, float8 should replace floatx.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Oh I left it in since it's still in use in other parts of the code base (autoquant, autoquant v2), and I wasn't sure if I should be touching those - is it ok to replace all instances across the whole codebase?
Yes, we want all the instances replaced. Autoquant is using it for Float8. Hence would be better to rename it float8 |
1 similar comment
This comment has been minimized.
This comment has been minimized.
Done! |
@@ -209,19 +215,64 @@ def __repr__(self): | |||
) | |||
|
|||
|
|||
class Float8QuantizedTensor(AffineQuantizedTensor): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I'm not a fan of this, this introduces one more abstraction (Float8QuantizedTensor
), while keeping the complexity of AffineQuantizedTensor
. I think either staying with AQT or just writing a float8 tensor without using AQT would seem more attractive.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Interesting - cc @jainapurva @jerryzh168 thoughts on this?
For context AQT subclassing was part of a BE effort for the week, I'll share the doc with you internally
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Removing AQT abstraction is easy, but the only reason I felt like keeping it was consistency in all dtypes. Though I do agree that it adds another level of abstraction
Discussed offline, closing until internal discussions are finalized. |
Context
Currently, AQT has the method from_hp_to_floatx for float8 quantization, and from_hp_to_fpx for low precision floating point data types like fp6 (technically can support fp1-fp7).
from_hp_to_floatx
re-uses from_hp_to_intx, which in turn uses these generic quantization primitives.Overall, in the current state the float8 path is a bit confusing for developers, due to both the naming ("floatx") and the use of generic functions which include a bunch of params which are unrelated to float8 quantization.
Summary of changes
The goal of this PR stack is to refactor this to have a clean separation of concerns, and simpler internal API surfaces for code using in float8 quantization for inference.
Specifically:
Note: I will add float8 static quantization in a separate set of PRs.