[Reland][CPU] Support int8 scaled embedding bag #3060
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/3060
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 02a1bc4 with merge base 0d3217d ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
bot
added
the
CLA Signed
Xia-Weiwen
changed the title
Xia-Weiwen
added
the
topic: not user facing
shiyang-weng
force-pushed
the
wengshiy/int8_scaled_embedding_bag
branch
from
ac75d21 to
e846df7
Compare
| } | ||
|
|
||
| } // namespace torchao | ||
| } // namespace torchao |
|
CC @mingfeima for review. Thanks. |
| # Next setp: support more out_dtype | ||
| out_dtype = torch.float32 |
There was a problem hiding this comment.
should this arg be exposed to the op as well in the future?
There was a problem hiding this comment.
should this arg be exposed to the op as well in the future?
Yes
|
Hi @mingfeima Could you please review this PR? Thanks. |
| #endif | ||
|
|
||
| template <typename index_t> | ||
| template <typename index_t, typename data_t> |
There was a problem hiding this comment.
this is not mandatory, but you can use block_dim as an template argument and make this function simpler.
template <typename index_t, typename scalar_t, int block_dim>
|
one more thing, it would be better to refactor the fp8 conversition simd code with https://github.com/pytorch/ao/blob/main/torchao/csrc/cpu/aten_kernels/float8_linear.cpp, maybe put them in to vec util.h. |
4 participants
int8 scaled_embedding_bag reverted by #2974
On #2972, they shared error reason related to the unused qtype.
We re-enable int8 scaled_embedding_bag and fix the unused variables issue on this PR