[oneDNN] lookup_table op with support for BF16 data type. #31558
Conversation
|
@jczaja @wozna @wojtuss @arlesniak @lidanqing-intel Please start your review. |
python/paddle/fluid/tests/unittests/test_lookup_table_bf16_op.py
Outdated
Show resolved
Hide resolved
python/paddle/fluid/tests/unittests/test_lookup_table_bf16_op.py
Outdated
Show resolved
Hide resolved
| @@ -102,7 +102,8 @@ class LookupTableKernel : public framework::OpKernel<T> { | |||
| auto id_index = table_t.GetIndexFromId(ids[i]); | |||
|
|
|||
| if (id_index != -1) { | |||
| if (input_data_type == framework::proto::VarType::INT8) { | |||
| if (input_data_type == framework::proto::VarType::INT8 || | |||
| input_data_type == framework::proto::VarType::BF16) { | |||
There was a problem hiding this comment.
I wonder why do you change lookup_table_op but not lookup_table_v2_op?
There was a problem hiding this comment.
That's because in the model "word2vec" there is lookup_table op used. Moreover paddle.fluid.layers.embedding Py API creates lookup_table op.
| @@ -33,10 +33,19 @@ | |||
| from paddle.fluid.op import Operator | |||
| from paddle.fluid.executor import Executor | |||
| from paddle.fluid.framework import Program, OpProtoHolder, Variable | |||
| from testsuite import create_op, set_input, append_input_output, append_loss_ops | |||
| from paddle.fluid.tests.unittests.testsuite import ( | |||
There was a problem hiding this comment.
Since this is the recommended Python style for importing multiple modules from single package PEP328 and generally it looks more readable. Moreover this way you can have multiline import statement.
|
@luotao1 regarding issues rised in PR-CI-APPROVAL:
|
PR types
New features
PR changes
OPs
Describe
This PR adds support for BF16 data type in
lookup_tableop for both forward and backward pass.