Skip to content
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.

Sign up for GitHub

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

Jump to bottom

integrate new float8 quantization primitives into AQT #1598

Open
wants to merge 11 commits into
base: gh/danielvegamyhre/23/head
Choose a base branch
from
Open

integrate new float8 quantization primitives into AQT #1598

wants to merge 11 commits into from

Conversation

danielvegamyhre
Copy link
Contributor

@danielvegamyhre danielvegamyhre commented Jan 22, 2025 •
edited
Loading

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:

  • Separate quantization primitives for float8
  • Integrate those new quant primitives into AQT <------------------- (this PR)

Note: I will add float8 static quantization in a separate set of PRs.

@danielvegamyhre
Update
ab8d5b5 
[ghstack-poisoned]
@danielvegamyhre
Copy link
Contributor Author

danielvegamyhre commented Jan 22, 2025 •
edited
Loading

Stack from ghstack (oldest at bottom):

@pytorch-bot PyTorch Bot
Copy link

pytorch-bot bot commented Jan 22, 2025 •
edited
Loading

🔗 Helpful Links

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

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

✅ No Failures

As of commit 2cac42e with merge base 860da26 (image):
💚 Looks good so far! There are no failures yet. 💚

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

@danielvegamyhre danielvegamyhre mentioned this pull request Jan 22, 2025
Merged
@facebook-github-bot 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 Jan 22, 2025
danielvegamyhre added a commit that referenced this pull request Jan 22, 2025
@danielvegamyhre
integrate new float8 quantization primitives into AQT
65ceb64 
ghstack-source-id: 9aacd39f35698b23f70c6270aa67e95aff24e29a
ghstack-comment-id: 2608090492
Pull Request resolved: #1598
@danielvegamyhre danielvegamyhre added quantize topic: improvement Use this tag if this PR is an improvement (doesn't fit into any of the other categories) labels Jan 22, 2025
@danielvegamyhre danielvegamyhre mentioned this pull request Jan 22, 2025
Closed
@danielvegamyhre danielvegamyhre requested review from jainapurva and jerryzh168 and removed request for jainapurva January 22, 2025 19:36
vkuzo
vkuzo reviewed Jan 23, 2025
View reviewed changes
torchao/dtypes/affine_quantized_tensor.py Outdated
target_dtype,
)
fp8_data = _layout.post_process(fp8_data)
tensor_impl_ctr = get_tensor_impl_constructor(type(_layout))
Copy link
Contributor

Choose a reason for hiding this comment

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

does this have multiple options for float8? if not, just call it directly to reduce # of abstractions the code reader needs to know about

Copy link
Contributor Author

Choose a reason for hiding this comment

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

good point, done!

vkuzo
vkuzo reviewed Jan 23, 2025
View reviewed changes
torchao/dtypes/affine_quantized_tensor.py
target_dtype: torch.dtype,
block_size: Tuple[int, ...],
_layout: Layout = PlainLayout(),
):
Copy link
Contributor

Choose a reason for hiding this comment

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

a docblock here should explain the difference between from_hp_to_floatx, from_hp_to_fpx, from_hp_to_float8

danielvegamyhre added a commit that referenced this pull request Jan 23, 2025
@danielvegamyhre
integrate new float8 quantization primitives into AQT
cf84d4a 
ghstack-source-id: 9aacd39f35698b23f70c6270aa67e95aff24e29a
ghstack-comment-id: 2608090492
Pull Request resolved: #1598
@danielvegamyhre
Update
76c7bde 
[ghstack-poisoned]
danielvegamyhre added a commit that referenced this pull request Jan 23, 2025
@danielvegamyhre
integrate new float8 quantization primitives into AQT
02e82f3 
ghstack-source-id: c87842f417c0a4099e6667d7d27040c5104cec66
ghstack-comment-id: 2608090492
Pull Request resolved: #1598
jainapurva
jainapurva reviewed Jan 23, 2025
View reviewed changes
torchao/dtypes/affine_quantized_tensor.py
@@ -422,6 +417,39 @@ def from_hp_to_fpx(
tensor_impl = tensor_impl_ctr(floatx_packed, scale, None, _layout)
return cls(tensor_impl, block_size, original_shape, dtype=input_float.dtype)

@classmethod
def from_hp_to_float8(
Copy link
Contributor

Choose a reason for hiding this comment

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

Update from_hp_to_floatx with the new float8 logic. For fp1-fp7, we're using from_hp_to_floatx.

@danielvegamyhre
Update
e116ed8 
[ghstack-poisoned]
danielvegamyhre added a commit that referenced this pull request Jan 23, 2025
@danielvegamyhre
integrate new float8 quantization primitives into AQT
7cc4944 
ghstack-source-id: c1deeeb84bdbb109a245b7e6c84150a724a012e7
ghstack-comment-id: 2608090492
Pull Request resolved: #1598
@danielvegamyhre
Update
2888180 
[ghstack-poisoned]
danielvegamyhre added a commit that referenced this pull request Jan 23, 2025
@danielvegamyhre
integrate new float8 quantization primitives into AQT
a4cfe7d 
ghstack-source-id: c1deeeb84bdbb109a245b7e6c84150a724a012e7
ghstack-comment-id: 2608090492
Pull Request resolved: #1598
danielvegamyhre added a commit that referenced this pull request Jan 23, 2025
@danielvegamyhre
integrate new float8 quantization primitives into AQT
17f0d20 
ghstack-source-id: c1deeeb84bdbb109a245b7e6c84150a724a012e7
ghstack-comment-id: 2608090492
Pull Request resolved: #1598
@danielvegamyhre
Update
6ba7924 
[ghstack-poisoned]
danielvegamyhre added a commit that referenced this pull request Jan 24, 2025
@danielvegamyhre
integrate new float8 quantization primitives into AQT
be2d094 
ghstack-source-id: 982ea07cc649320a57bf47120489bac86f2e900b
ghstack-comment-id: 2608090492
Pull Request resolved: #1598
@danielvegamyhre
Update
7cd5bc4 
[ghstack-poisoned]
danielvegamyhre added a commit that referenced this pull request Jan 24, 2025
@danielvegamyhre
integrate new float8 quantization primitives into AQT
f419cac 
ghstack-source-id: e0a6b7977a147513dcac7ae8dc5ef6008f51847d
ghstack-comment-id: 2608090492
Pull Request resolved: #1598
@danielvegamyhre
Update
f8eae87 
[ghstack-poisoned]
danielvegamyhre added a commit that referenced this pull request Jan 24, 2025
@danielvegamyhre
integrate new float8 quantization primitives into AQT
0dbe010 
ghstack-source-id: fccae983561f2c5cb5c9ab6bd1b09e761e12ae23
ghstack-comment-id: 2608090492
Pull Request resolved: #1598
@danielvegamyhre
Update
9444e47 
[ghstack-poisoned]
danielvegamyhre added a commit that referenced this pull request Jan 24, 2025
@danielvegamyhre
integrate new float8 quantization primitives into AQT
afc2779 
ghstack-source-id: 2d1d4f154bd1ae9be39996695c4fe99c6768a8c5
ghstack-comment-id: 2608090492
Pull Request resolved: #1598
@danielvegamyhre
Update
e614e66 
[ghstack-poisoned]
danielvegamyhre added a commit that referenced this pull request Jan 24, 2025
@danielvegamyhre
integrate new float8 quantization primitives into AQT
41357fa 
ghstack-source-id: 0a9fcdb46e5ed4ea33f4a76dece6be02b4976bbe
ghstack-comment-id: 2608090492
Pull Request resolved: #1598
@danielvegamyhre
Update
20476f9 
[ghstack-poisoned]
danielvegamyhre added a commit that referenced this pull request Jan 24, 2025
@danielvegamyhre
integrate new float8 quantization primitives into AQT
9a1c85a 
ghstack-source-id: a098c0ec8afde88cc8839ca5c5e16a33089b208f
ghstack-comment-id: 2608090492
Pull Request resolved: #1598
@danielvegamyhre
Update
2cac42e 
[ghstack-poisoned]
danielvegamyhre added a commit that referenced this pull request Jan 24, 2025
@danielvegamyhre
integrate new float8 quantization primitives into AQT
d24b70e 
ghstack-source-id: b9139dc6b848d70a209b3a9736a73fe37edef9a2
ghstack-comment-id: 2608090492
Pull Request resolved: #1598
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@jainapurva jainapurva jainapurva left review comments

@vkuzo vkuzo vkuzo left review comments

@jerryzh168 jerryzh168 Awaiting requested review from jerryzh168

Assignees
No one assigned
Labels
CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. quantize topic: improvement Use this tag if this PR is an improvement (doesn't fit into any of the other categories)
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
4 participants
@danielvegamyhre @vkuzo @jainapurva @facebook-github-bot