New embedding quant fusion #10325
New embedding quant fusion #10325
Conversation
metascroy
requested review from
JacobSzwejbka,
tarun292 and
larryliu0820
as code owners
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/10325
Note: Links to docs will display an error until the docs builds have been completed. ❌ 2 New FailuresAs of commit d88601f with merge base ad1b154 ( NEW FAILURES - The following jobs have failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
facebook-github-bot
added
the
CLA Signed
|
This pull request was exported from Phabricator. Differential Revision: D73381542 |
|
@metascroy has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
| weight_1 = weight_view[:, :, 1] << 2 | ||
| weight_2 = weight_view[:, :, 2] << 4 | ||
| weight_3 = weight_view[:, :, 3] << 6 | ||
| packed_weight = weight_0 + weight_1 + weight_2 + weight_3 |
There was a problem hiding this comment.
nit: Just do bitwise OR.
| weight_view = weight_range_shifted.view( | ||
| weight.shape[0], weight.shape[1] // 2, 2 | ||
| ) | ||
| weight_even = weight_view[:, :, 0] * 16 # left shift 4 |
There was a problem hiding this comment.
Why * 16 here but shift by 4 above?
There was a problem hiding this comment.
Isn't * 16 the same as shift by 4?
I just copied the code from here: https://github.com/pytorch/executorch/blob/main/examples/models/llama/source_transformation/quantize.py#L659-L683
But I can clean it up and change to using left shift.
There was a problem hiding this comment.
yeah was just highlighting for consistency
| ) | ||
| weight_even = weight_view[:, :, 0] * 16 # left shift 4 | ||
| weight_odd = weight_view[:, :, 1] | ||
| packed_weight = weight_even + weight_odd |
| @@ -296,6 +362,22 @@ def embedding_2bit_dtype( | |||
| return torch.ops.aten.embedding.default(weight, indices) | |||
|
|
|||
|
|
|||
| @register_fake("quantized_decomposed::embedding_2bit.dtype") | |||
There was a problem hiding this comment.
I think so? You cannot dynamo trace code that doesn't have meta kernels registered.
There was a problem hiding this comment.
yeah sorry I think I asked because I wasnt sure how it was working before without it. But yeah you do need meta kernel
| num_embeddings, packed_embedding_dim = weight.shape | ||
| embedding_dim = packed_embedding_dim * 2 | ||
| embedding = torch.nn.Embedding(num_embeddings, embedding_dim, device=weight.device) |
There was a problem hiding this comment.
very small nit: small refactor can abstract these 3 lines out.
| def embedding_byte_dtype_pattern( | ||
| indices, int_data, group_size, scale, zero_point, output_dtype | ||
| ): | ||
| dq = torch.ops.torchao.dequantize_affine.default( |
There was a problem hiding this comment.
what does "INT" mean here?
There was a problem hiding this comment.
It's the zero point domain. It means the zero points are integers. But this is an arg torchao is going to change in their quant primitives when they clean them up.
| return torch.ops.aten.embedding.default(dq, indices) | ||
|
|
||
| def embedding_byte_replacement(indices, int_data, group_size, scale, zero_point): | ||
| zero_point_dtype_cast = torch.ops.aten.to.dtype(zero_point, scale.dtype) |
There was a problem hiding this comment.
embedding byte ops take float zero point?
There was a problem hiding this comment.
The ones in the ET kernels do: https://github.com/pytorch/executorch/blob/main/kernels/quantized/cpu/embeddingxb.cpp#L200
| @@ -298,6 +298,8 @@ python_unittest( | |||
| "//caffe2:torch", | |||
| "//executorch/exir:lib", | |||
| "//executorch/exir/passes:quant_fusion_pass", | |||
| "//pytorch/ao:torchao", | |||
| "//executorch/exir/passes:constant_prop_pass", | |||
There was a problem hiding this comment.
shouldnt you call const prop in the QuantFusionPass?
There was a problem hiding this comment.
Unfortunately QuantFusionPass works on a graph module and const_prop_pass works on an exported_program. I thought we could enable both in to_executorch by default, rather than have users call the passes separately like is done in the unit test?
There was a problem hiding this comment.
I think const prop has some nuances which may make it non-trivial for excample for q/dq nodes.
My concern here would be the perf cliff for the uninitiated.
Can we just not const prop manually? instead of even introducing this op in the graph in the first place?
There was a problem hiding this comment.
Manual const propagation would still require updating the signature of the exported program because we're changing the weights. So it couldn't be done on the graph module.
In terms of perf cliff, it will not lower to ExecuTorch without const propagation because the pack embedding op has no out-variant.
There was a problem hiding this comment.
I see. ok. but lets plan to have this cost prop done appropriately. I highly doubt that this can be done transparently. For example for quantized models' have dq on weights that might get const propagated
exir/tests/test_quant_fusion_pass.py
Outdated
| self._test_embedding_torchao(bit_width, test_dtype_variant, test_per_group) | ||
|
|
||
| def _test_embedding_torchao( | ||
| self, bit_width: int, test_dtype_variant: bool, test_per_group: bool |
There was a problem hiding this comment.
test_dtype_variant as bool feels like a bit of misnomer
There was a problem hiding this comment.
I'll change to use_dtype_variant?
There was a problem hiding this comment.
i think i misunderstood this probably because it was a bit late when i was reviewin git. You just meant to check .dtype variant of the op and I thought you were checking fp16 vs fp32 vs other dtype variants
Summary:
The diff adds new quant fusion passes to recognize 2, 4, and 8 bit quantized embeedings (per group and per channel) and fuses them to ExecuTorch kernels. This makes torchao's quantize_ integrate with ExecuTorch:
```
quantize_(
model,
IntxWeightOnlyConfig(weight_dtype=torch.int4, granularity=PerGroup(32)),
lambda m, fqn: isinstance(m, torch.nn.Embedding)
)
# lower model to executorch
```
For the model to lower, we need to run QuantFusionPass. For subbyte, we also need to run constant_prop_pass. (See new unit tests for examples). In follow-up diffs, we will enable these passes by default in to_executorch before the memory passing and out-variant passes.
Differential Revision: D73381542
metascroy
force-pushed
the
export-D73381542
branch
from
ff15175 to
d88601f
Compare
|
@metascroy has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
Summary:
The diff adds new quant fusion passes to recognize 2, 4, and 8 bit quantized embeedings (per group and per channel) and fuses them to ExecuTorch kernels. This makes torchao's quantize_ integrate with ExecuTorch:
For the model to lower, we need to run QuantFusionPass. For subbyte, we also need to run constant_prop_pass. (See new unit tests for examples). In follow-up diffs, we will enable these passes by default in to_executorch before the memory passing and out-variant passes.
Differential Revision: D73381542