Exported TF Bert model is much slower than that exported from Google's Bert

[jlei2]

I benchmarked the Bert model exported from Huggingface TF Bert code and observed that it is much slower than Google's TF2 Bert model(https://github.com/tensorflow/models/tree/master/official/nlp/bert) on TF Serving (GPU device). When batch_size=1, Huggingface Bert model can be more than 200% slower than Google's one. Here are my benchmark results:

latency: (ms/batch) on 1 GPU-V100

Batch Size	google official	HF_TF2	relative change
1	6.7	21.3	238.10%
2	9.4	24.2	171.91%
4	14.4	28.1	99.29%
8	24.0	36.9	55.70%
16	46.6	58.6	26.57%
64	171.5	171.4	0.35%
128	338.5	324.5	-3.71%

I used the latest Tensorflow profiling tool (https://www.tensorflow.org/tfx/serving/tensorboard) to compare these two:

Hugging-face Bert-Base FP32 Batch_size = 128

Google Bert-Base FP32 Batch_size = 128

Comparing the event traces of HF's and google's, we see that before each MatMul operation on GPU, there are always some CPU processes running (wasn’t able to get more info about these CPU processes because the Profiling tool doesn’t provide meaningful labels or description for those CPU processes but it seems like they are doing some gather/reduce operations) and the GPU is inactive for a while. If you look at the trace of google official, you can find that all GPU ops are compact and there are no idle time for GPU. In parallel, only one CPU process keeps running and doesn’t cause GPU ops to wait. The short idle time will hurt the performance more when batch_size or model_size is small because short time is spent on GPU calculations and idle time is not ignorable anymore. See time comparison grouped by type below:

Hugging-face Bert-Base FP32 Batch_size = 4

Google Bert-Base FP32 Batch_size = 4

Under small batch_size=4, Hugging-face model has 60% GPU idle time while that of google model is only 25%. This also explains the reason why for small batch_size, there is over 200% slow down and for large batch_size=128, the situation gets better.

The MatMul op is triggered by tf.keras.layers.Dense(), which is widely used in Transformer encoder self-attention, intermediate layer and output layer. In comparison, Google's Bert uses DenseEinsum()(link) to replace all usage of Dense layer. I personally modified Huggingface's TF Bert code base to use DenseEinsum(). And the slow down issue got solved:

latency: (ms/batch) on 1 GPU-V100

Batch Size	google official	HF_TF2	relative change	After Fixing	relative change
1	6.7	21.3	238.10%	6.60	4.76%
2	9.4	24.2	171.91%	8.90	0.00%
4	14.4	28.1	99.29%	13.40	-4.96%
8	24.0	36.9	55.70%	22.10	-6.75%
16	46.6	58.6	26.57%	43.20	-6.70%
64	171.5	171.4	0.35%	158.80	-7.03%
128	338.5	324.5	-3.71%	313.10	-7.09%

Profile after fixing:

I noticed that multiple issues may be related to this (#6264).

Do you plan to solve this issue? Thanks. @patrickvonplaten

[patrickvonplaten]

This looks super interesting @jlei2 ! I openen a PR #6877 to see if changing to EinusmDense speeds up the runtime...In the PR I cannot see an improvement - could you take a look and maybe comment on how you changed the TF Bert code in transformers to benchmark your changes? :-)

[jlei2]

Hi Patrick, thank you for taking time to investigate this issue! Your PR is not completely correct. Here is my draft PR. DenseEinsum actually is the combination of matrix reshape/transpose and dense layer. So we need to be careful when replacing keras.dense with DenseEinsum, otherwise we will not be able to get the same output as we can get before this change.

I tested on my end that there is no speedup after this change at runtime, using the huggingface bechmark command tool given by you. But I did see expected speedup after exporting DenseEinsum models. I used tensorflow profiling tool(https://www.tensorflow.org/tfx/serving/tensorboard) to benchmark the inference time of exported model when serving on TF-serving 2.2.0. So actually I don't have my own benchmark code.

Device: 1 GPU-V100
Batch Size: 1
Model: Bert-base FP32

So the left picture is from the Huggingface model after applying my PR. The right one is from original Huggingface model using current master. You can see that there is almost 100% speedup. So I suspect this issue only happens on exported model for tf-serving.

The code I used to export saved model:

import tensorflow as tf
import transformers

from transformers import BertModel
pt_model_path = 'bert-base-uncased'

model = BertModel.from_pretrained('bert-base-uncased')
model.save_pretrained(pt_model_path)

MAX_LEN = 128
model_path = 'saved_model/tmp_model'
input_ids = tf.keras.layers.Input((MAX_LEN,), dtype=tf.int32, name='input_ids')
attention_mask = tf.keras.layers.Input((MAX_LEN,), dtype=tf.int32, name='attention_mask')
token_type_ids = tf.keras.layers.Input((MAX_LEN,), dtype=tf.int32, name='token_type_ids')
base_model = transformers.TFBertModel.from_pretrained(pt_model_path, output_hidden_states=False, from_pt=True)
base_output = base_model.bert([input_ids, attention_mask, token_type_ids])
seq_out, pool_out = base_output[0], base_output[1]
base_model.trainable = False
model = tf.keras.models.Model(inputs=[input_ids, attention_mask, token_type_ids], outputs=[pool_out])
model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
print(model.summary())
model.save(model_path, include_optimizer=False, save_format="tf")

Is this something you will want to solve on your end? though we see that during runtime there is no difference.

[patrickvonplaten]

Hey @jlei2,

Thanks a lot for checking and your draft PR. I think we will be able to implement your proposition given that we can keep complete backwards compatibility. Do you have an idea by any chance why we can see the speed up only on TF-serving?

Also cc @jplu - this looks very interesting!

[jlei2]

To be honest I have no idea about the root difference on TFS side. I can only observe that there are some CPU processes wasting time before the Matmul op. My feeling is that the MatMul op triggered by tf.keras.layers.Dense() may not be implemented to be very efficient on TF-serving. Though this issue seems to have more to do with tensorflow keras team or TFS team instead of Huggingface code base imo, it would be helpful to all Huggingface users if you are able to resolve this issue on your end.

[stale]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.

[thomwolf]

Re-activated cc @jplu

[jplu]

Hey @jlei2!

Sorry for this long silence. We are currently working on this performance issue and to be sure to make the proper changes, can you share with us your benchmark script that you used to compute the values in the table?

Thanks!