Add per tensor fp8 conv2d support #3315
Merged
+218
−46
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/3315
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (9 Unrelated Failures)As of commit 22d7227 with merge base e8c4d09 ( BROKEN TRUNK - The following jobs failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
bot
added
the
CLA Signed
jerryzh168
force-pushed
the
add-conv2d
branch
from
007ea02 to
8bf4032
Compare
jerryzh168
added
the
topic: new feature
andrewor14
approved these changes
Nov 10, 2025
| padding = [0, *padding] | ||
| stride = [1, *stride] | ||
| dilation = [1, *dilation] | ||
| res = _quantize_and_scaled_conv3d( |
Contributor
There was a problem hiding this comment.
is this unsqueezing to turn 2d into 3d? Does fbgemm only have 3d conv kernels?
Contributor
Author
There was a problem hiding this comment.
yeah this turns 2d conv to 3d and it's because fbgemm only supports 3d conv right now
|
|
||
| if weight.dim() == 5: | ||
| # weights for conv3d | ||
| if weight.dim() in [4, 5]: |
Contributor
There was a problem hiding this comment.
is there anything more robust we can check here?
Contributor
Author
There was a problem hiding this comment.
we can't really distinguish from here whether it's linear and conv weight I think (although right now seems linear is 2d/3d and conv is 4d/5d, maybe conv1d could have 3d weight, which is an overlap with linear)
but we could potentially separate the handling of conv and linear by passing around the module as well, if this is needed in the future
I can add a comment for now
| "Please make sure both activation and weights are in the `channels_last` memory_format" | ||
| ) | ||
| input_tensor = input_tensor.unsqueeze(2) | ||
| weight_tensor = weight_tensor.unsqueeze(2) |
Contributor
There was a problem hiding this comment.
Can we reuse this code, e.g. have aten.conv2d call into aten.convolution?
Contributor
Author
There was a problem hiding this comment.
yeah I think that should be possible for both aten.conv2d and aten.conv3d, I can refactor that in next PR?
jerryzh168
force-pushed
the
add-conv2d
branch
2 times, most recently
from
fe50ba6 to
9cf37be
Compare
Summary: Add fp8 conv2d support, using the same conv3d kernels, by setting the D dimension to 1. 1. unsqueeze both input and weight in dim 2 ( the D dimension) 2. call fp8 conv3d op from fbgemm `torch.ops.fbgemm.f8f8bf16_conv` 3. assert D dimension shape to be 1 and call sequeeze at dim 2: res.squeeze(2) to remove the D dimension Test Plan: python test/quantization/quantize_/workflows/float8/test_float8_tensor.py -k test_unsqueeze_conv2d_weight python test/quantization/quantize_/workflows/float8/test_float8_tensor.py -k test_fp8_conv_variants
jerryzh168
force-pushed
the
add-conv2d
branch
from
9cf37be to
22d7227
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
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.
Summary:
Add fp8 conv2d support, using the same conv3d kernels, by setting the D dimension to 1.
torch.ops.fbgemm.f8f8bf16_convTest Plan: