Update coreml codebook #2648
Merged
Update coreml codebook #2648
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/2648
Note: Links to docs will display an error until the docs builds have been completed. ❌ 2 New Failures, 1 PendingAs of commit cceaab2 with merge base 22f9d31 ( NEW FAILURES - The following jobs have failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
bot
added
the
CLA Signed
This was referenced Jul 31, 2025
metascroy
added
topic: not user facing
jerryzh168
reviewed
Aug 1, 2025
| weight_tensor = weight_tensor.dequantize() | ||
| return func(indices, weight_tensor, **kwargs) | ||
|
|
||
|
|
||
| @implements([aten.detach.default, aten.alias.default]) |
There was a problem hiding this comment.
This I will save for a future PR
jerryzh168
reviewed
Aug 1, 2025
| @@ -48,43 +49,45 @@ def choose_qparams_and_quantize_codebook_coreml( | |||
|
|
|||
| Returns: | |||
| Tuple[torch.Tensor, torch.Tensor] The codebook (lookup table) Tensor and the quantized Tensor (codes, torch.uint8) | |||
| The LUT table has dimension input_tensor.dim() + 2, where: | |||
There was a problem hiding this comment.
nit: maybe spell out the dimensions with variables/expressions?
jerryzh168
reviewed
Aug 1, 2025
| dequant = torch.zeros_like(codes, dtype=output_dtype) | ||
| # Compute shape of lookup group indices from codes shape and block size | ||
| code_shape = codes.shape | ||
| ndim = len(code_shape) |
jerryzh168
reviewed
Aug 1, 2025
| dequant[:, i, :] = codebook[i][codes[:, i, :]].squeeze() | ||
| # Compute which codebook slice to use for each element | ||
| group_indices = [] | ||
| for dim, bsz in zip(code_shape, block_size): |
There was a problem hiding this comment.
nit: dim might be confusing? this seems to be code_size_i and block_size_i
jerryzh168
reviewed
Aug 1, 2025
Comment on lines
188
to
197
| mesh = torch.meshgrid(*group_indices, indexing="ij") | ||
| group_index_tensor = torch.stack(mesh, dim=-1) # shape (..., N), where N = ndim | ||
|
|
||
| # Flatten everything to index efficiently | ||
| flat_codes = codes.reshape(-1) | ||
| flat_groups = group_index_tensor.reshape(-1, ndim) # (..., ndim) | ||
|
|
||
| # Compute dequantized values via indexing | ||
| # index into codebook with (*group_index, code_index, :) | ||
| gathered = codebook[(*flat_groups.T, flat_codes)] # shape (numel, vec_dim) |
There was a problem hiding this comment.
nit: can you add comments of some examples for these to make it easier to understand
jerryzh168
approved these changes
Aug 1, 2025
metascroy
added a commit
to pytorch/executorch
that referenced
this pull request
Aug 6, 2025
This adds palletization support for embedding/linear layers in CoreML using TorchAO's quantize_ API. Note, this needs to wait for pytorch/ao#2648 to land in ao + a pin bump in ET before landing.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This PR updates the dequantize_codebook quant primitive to be more compatible with CoreML. More specifically:
code_dtype is changed to nbits because CoreML cannot process a function that has non-standard dtypes in the signature (e.g., torch.uint3)
This changes the codebook rank to codes.dim() + 2 to follow the convention here https://apple.github.io/coremltools/source/coremltools.converters.mil.mil.ops.defs.html#coremltools.converters.mil.mil.ops.defs.iOS18.compression.constexpr_lut_to_dense