Skip to content

Enable PerRow(axis) to support axes other than -1 #3303

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
Open
wants to merge 13 commits into
base: main
Choose a base branch
from
Open

Enable PerRow(axis) to support axes other than -1 #3303

wants to merge 13 commits into from

Conversation

@vkuzo
Copy link
Contributor

@vkuzo vkuzo commented Nov 6, 2025 •
edited
Loading

Summary:

Enables per-row float8 quantization of tensors with shape (B, K, N) across the K dimension.

TODO before land:

  1. [done] audit PerRow callsites across the codebase and see if anything else needs changing
  2. [done] verify serialization will not break

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

vkuzo added 3 commits November 5, 2025 12:49
@vkuzo
Update
d76d3ff 
[ghstack-poisoned]
@vkuzo
Update
c97b030 
[ghstack-poisoned]
@vkuzo
Update
763a709 
[ghstack-poisoned]
@vkuzo
Copy link
Contributor Author

vkuzo commented Nov 6, 2025 •
edited
Loading

Stack from ghstack (oldest at bottom):

@pytorch-bot
Copy link

pytorch-bot bot commented Nov 6, 2025 •
edited
Loading

🔗 Helpful Links

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

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

✅ No Failures

As of commit b0d38ad with merge base 6259e98 (image):
💚 Looks good so far! There are no failures yet. 💚

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

vkuzo added a commit that referenced this pull request Nov 6, 2025
@vkuzo
[wip] float8 rowwise quant along row 1 of tensor rank 2
07eb188 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
ghstack-source-id: c40d96b
ghstack-comment-id: 3497584430
Pull-Request: #3303
@meta-cla meta-cla 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 6, 2025
@vkuzo vkuzo mentioned this pull request Nov 6, 2025
Merged
@vkuzo
Update
6143910 
[ghstack-poisoned]
vkuzo added a commit that referenced this pull request Nov 6, 2025
@vkuzo
[wip] float8 rowwise quant along row 1 of tensor rank 2
83c239e 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
ghstack-source-id: a113126
ghstack-comment-id: 3497584430
Pull-Request: #3303
@vkuzo vkuzo changed the title [wip] float8 rowwise quant along row 1 of tensor rank 2 Enable PerRow(axis) to support axes other than -1 Nov 6, 2025
@vkuzo vkuzo added the topic: improvement Use this tag if this PR is an improvement (doesn't fit into any of the other categories) label Nov 6, 2025
jcaip
jcaip reviewed Nov 6, 2025
View reviewed changes
jerryzh168
jerryzh168 reviewed Nov 6, 2025
View reviewed changes
@vkuzo
Update
0073d60 
[ghstack-poisoned]
vkuzo added a commit that referenced this pull request Nov 6, 2025
@vkuzo
[wip] float8 rowwise quant along row 1 of tensor rank 2
bc07b7b 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
ghstack-source-id: 688ac7b
ghstack-comment-id: 3497584430
Pull-Request: #3303
@vkuzo
Update
d9ebeb3 
[ghstack-poisoned]
vkuzo added a commit that referenced this pull request Nov 6, 2025
@vkuzo
[wip] float8 rowwise quant along row 1 of tensor rank 2
03822c8 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
ghstack-source-id: efb2efb
ghstack-comment-id: 3497584430
Pull-Request: #3303
@vkuzo vkuzo changed the base branch from gh/vkuzo/164/head to main November 6, 2025 21:32
@vkuzo
Update
a99196f 
[ghstack-poisoned]
vkuzo added a commit that referenced this pull request Nov 6, 2025
@vkuzo
[wip] float8 rowwise quant along row 1 of tensor rank 2
dc92bb1 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
ghstack-source-id: d3d5c06
ghstack-comment-id: 3497584430
Pull-Request: #3303
jerryzh168
jerryzh168 approved these changes Nov 6, 2025
View reviewed changes
Copy link
Contributor

@jerryzh168 jerryzh168 left a comment

Choose a reason for hiding this comment

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

looks good

@vkuzo
Update
84b06d1 
[ghstack-poisoned]
vkuzo added a commit that referenced this pull request Nov 7, 2025
@vkuzo
[wip] float8 rowwise quant along row 1 of tensor rank 2
f383884 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
ghstack-source-id: 1a10e81
ghstack-comment-id: 3497584430
Pull-Request: #3303
@vkuzo
Update
d4ce355 
[ghstack-poisoned]
vkuzo added a commit that referenced this pull request Nov 7, 2025
@vkuzo
[wip] float8 rowwise quant along row 1 of tensor rank 2
6f8deff 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
ghstack-source-id: f396611
ghstack-comment-id: 3497584430
Pull-Request: #3303
@vkuzo
Update
b49280d 
[ghstack-poisoned]
vkuzo added a commit that referenced this pull request Nov 7, 2025
@vkuzo
[wip] float8 rowwise quant along row 1 of tensor rank 2
b2e0bcd 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
ghstack-source-id: f511df3
ghstack-comment-id: 3497584430
Pull-Request: #3303
vkuzo added a commit that referenced this pull request Nov 7, 2025
@vkuzo
[wip] float8 rowwise quant along row 1 of tensor rank 2
dfcad30 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
ghstack-source-id: f511df3
ghstack-comment-id: 3497584430
Pull-Request: #3303
@vkuzo
Update
ec28bf5 
[ghstack-poisoned]
vkuzo added a commit that referenced this pull request Nov 7, 2025
@vkuzo
[wip] float8 rowwise quant along row 1 of tensor rank 2
7b7892d 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
ghstack-source-id: 4c8849d
ghstack-comment-id: 3497584430
Pull-Request: #3303
vkuzo
vkuzo commented Nov 7, 2025
View reviewed changes
torchao/testing/utils.py
param_data = param.data
param_data = param_data.narrow(output_dim, start_idx, shard_size)
orig_value = param_data.qdata[0][0]
orig_values = param_data.qdata[0]
Copy link
Contributor Author

@vkuzo vkuzo Nov 7, 2025

Choose a reason for hiding this comment

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

the original check failed by chance, the two tensors had the same value at [0][0]. Checking the first row is more resistant to chance.

@vkuzo
Update
a6324b0 
[ghstack-poisoned]
vkuzo added a commit that referenced this pull request Nov 7, 2025
@vkuzo
[wip] float8 rowwise quant along row 1 of tensor rank 2
3f683dd 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
ghstack-source-id: c4363b4
ghstack-comment-id: 3497584430
Pull-Request: #3303
@vkuzo vkuzo mentioned this pull request Nov 7, 2025
Merged
@vkuzo
Update
b0d38ad 
[ghstack-poisoned]
vkuzo added a commit that referenced this pull request Nov 7, 2025
@vkuzo
[wip] float8 rowwise quant along row 1 of tensor rank 2
80bf726 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
ghstack-source-id: 901c271
ghstack-comment-id: 3497584430
Pull-Request: #3303
@jcaip jcaip mentioned this pull request Nov 7, 2025
Open
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@jcaip jcaip jcaip left review comments

@jerryzh168 jerryzh168 jerryzh168 approved these changes

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. 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

@vkuzo @jerryzh168 @jcaip