Skip to content

Adds Q/DQ layout support for embedding quantization with IntxWeightOnlyConfig#1972

Merged
metascroy merged 8 commits into
mainfrom
add-intx-weight-only-for-embedding
Apr 8, 2025
Merged

Adds Q/DQ layout support for embedding quantization with IntxWeightOnlyConfig#1972
metascroy merged 8 commits into
mainfrom
add-intx-weight-only-for-embedding

Conversation

@metascroy

Copy link
Copy Markdown
Contributor

This will be used to quantize embeddings in ET.

@pytorch-bot

pytorch-bot Bot commented Mar 27, 2025

Copy link
Copy Markdown

🔗 Helpful Links

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

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

❌ 1 New Failure

As of commit 186f903 with merge base 5ded23c (image):

NEW FAILURE - The following job has failed:

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

@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 Mar 27, 2025
@metascroy metascroy requested a review from jerryzh168 March 27, 2025 00:59
Comment thread torchao/experimental/quant_api.py Outdated
Comment thread torchao/experimental/quant_api.py Outdated
@metascroy metascroy force-pushed the add-intx-weight-only-for-embedding branch from 2c3b9ac to 05eec5d Compare April 6, 2025 02:15

@dataclass
class IntxWeightOnlyConfig(AOBaseConfig):
"""

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

@andrewor14 can you have a look at this comment if there are any issues with it working well with QAT workflow with FakeQuantizeConfig.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Strange that we have IntxWeightOnly and Int4WeightOnly

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

yeah I feel we should probably merge these two

Int8DynamicActivationIntxWeightConfig(
weight_dtype=weight_dtype,
granularity=granularity,
granularity=PerRow(),

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

should this be PerAxis as well

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

It can't be because that's controlled by Int8DynamicActivationIntxWeightConfig, which uses PerRow until #1968 lands

SharedEmbeddingQuantizer(
weight_dtype=weight_dtype,
granularity=granularity,
granularity=PerRow(),

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

and this

@jerryzh168 jerryzh168 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

looks good overall, just need to change PerRow to PerAxis(axis=0) as we discussed in meeting

@metascroy metascroy merged commit 9516764 into main Apr 8, 2025
jainapurva pushed a commit that referenced this pull request Apr 8, 2025
…lyConfig (#1972)

* up

* up

* up

* up

* up

* up

* up

* up
def _(func, types, args, kwargs):
if _embedding_q_dq_check(args, kwargs):
return _embedding_q_dq_impl(args, kwargs)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

why does line 299 only dequantizes weight bu tnot actually run embedding op?

liangel-02 pushed a commit that referenced this pull request Aug 25, 2025
…lyConfig (#1972)

* up

* up

* up

* up

* up

* up

* up

* up
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants