Exported TF Bert model is much slower than that exported from Google's Bert
Author: jlei2
Pull request: no

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**
<img width="1518" alt="Screen Shot 2020-06-05 at 8 20 34 PM" src="https://user-images.githubusercontent.com/70337521/91475490-87b47600-e850-11ea-8e19-ac2b5103073d.png">
**Google Bert-Base FP32 Batch_size = 128**
<img width="1518" alt="image (1)" src="https://user-images.githubusercontent.com/70337521/91475500-8a16d000-e850-11ea-880c-eab176c0db59.png">

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**  
<img width="425" alt="image (3)" src="https://user-images.githubusercontent.com/70337521/91475887-175a2480-e851-11ea-8493-a25f0cea4914.png">

**Google Bert-Base FP32 Batch_size = 4**
<img width="418" alt="image (4)" src="https://user-images.githubusercontent.com/70337521/91475882-145f3400-e851-11ea-830d-7ed59a04e3bc.png">

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](https://github.com/tensorflow/models/blob/1ebff962832ec8a599bd67ac8c43b9ae1b05c3fa/official/nlp/transformer/attention_layer.py#L59-L69)) 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:
<img width="1517" alt="image (4) copy" src="https://user-images.githubusercontent.com/70337521/91475978-3789e380-e851-11ea-9dae-a5136fac1074.png">

I noticed that multiple issues may be related to this (https://github.com/huggingface/transformers/issues/6264).

Do you plan to solve this issue? Thanks. @patrickvonplaten 


Comments:
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? :-) 

Hi Patrick, thank you for taking time to investigate this issue! Your PR is not completely correct. Here is my draft [PR](https://github.com/jlei2/transformers/pull/1). 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

<img width="220" alt="Screen Shot 2020-09-02 at 1 52 08 PM" src="https://user-images.githubusercontent.com/70337521/92035460-981e9200-ed23-11ea-8141-a5f8e600559a.png">  <img width="220" alt="Screen Shot 2020-09-02 at 1 52 28 PM" src="https://user-images.githubusercontent.com/70337521/92035489-a2409080-ed23-11ea-9414-dd2bb6e8ea68.png">
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.

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! 

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.

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.


Re-activated cc @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!

Hi @jplu , thanks for keeping an eye on this! Unfortunately I can not share the benchmark script I used to get the numbers in the table with you because it's confidential. But I can share how you can obtain the same slow-down observation using Tensorflow Profiler tool (to get the two Performance Summary Screenshots I uploaded above 15.7ms vs 28.8ms).

### Export baseline and modified SavedModel
**Baseline transformers repo**: The commit I used is the same as the master branch in my fork (https://github.com/jlei2/transformers)

**Modified transformers repo using tf.keras.layers.experimental.EinsumDense**: (https://github.com/jlei2/transformers/tree/dense-einsum-tf2.3)

**The code to export the 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")
```

Then on a machine with 1 GPU-V100 and TF-Serving-gpu== 2.3.0 and Tensorflow==2.3.0:
### Spin up SavedModel:
```
export LD_LIBRARY_PATH=/usr/local/cuda/extras/CUPTI/lib64:$LD_LIBRARY_PATH
export MODEL_DIR=your_model_dir
export MODEL_NAME=your_model_name
nohup tensorflow_model_server --rest_api_port=8501 --model_name=$MODEL_NAME --model_base_path=$MODEL_DIR >server.log 2>&1
```

### Spin up Tensorboard
```
tensorboard --logdir ~/logs/inference_demo/ --port your_port_number --bind_all
```

Go to the profile page on Tensorboard
<img width="447" alt="Screen Shot 2020-12-01 at 8 58 18 PM" src="https://user-images.githubusercontent.com/70337521/100830662-03f2c280-3419-11eb-95bc-f0d70caec6db.png">

After clicking capture, I would send request to TFS to profile the whole process

**Code to send request to served model**:
```
import json
import requests

BATCH_SIZE = 1
MAX_SEQ_LEN = 128
batch = []
MODEL_NAME = your_model_name

for _ in range(BATCH_SIZE):
    batch.append({"input_ids":[999] * MAX_SEQ_LEN, "attention_mask":[1] * MAX_SEQ_LEN, "token_type_ids":[0] * MAX_SEQ_LEN})

input_data = {"instances": batch}
r = requests.post("http://localhost:8501/v1/models/%s:predict"%(MODEL_NAME), data=json.dumps(input_data))
print(r.text)
```
I would run this scripts for several times first to warm up the model and then start to profile formally. 

And finally you will see profiling results on Tensorboard UI page just like what I uploaded.

Hope this could be helpful to you!

Actually I write a very simple benchmark script that can show the difference:

```
import json
import requests
import time
BATCH_SIZE = 1
MAX_SEQ_LEN = 128
batch = []
MODEL_NAME = your_model_name

for _ in range(BATCH_SIZE):
	batch.append({"input_ids":[999] * MAX_SEQ_LEN, "attention_mask":[1] * MAX_SEQ_LEN, "token_type_ids":[0] * MAX_SEQ_LEN})

input_data = {"instances": batch}

start = time.time()
for _ in range(100):
	r = requests.post("http://localhost:8501/v1/models/%s:predict"%(MODEL_NAME), data=json.dumps(input_data))
end = time.time()
print(end-start)
```
Baseline time: ~2.8s
My version's time: ~1.5s.

So we can easily see ~2x speed up.

Awesome!! Thanks a lot @jlei2!! This is already a good start to check the differences. I will let you know here once we have done the changes!

@jlei2 I have open a PR for integrating this change. Unfortunately, as I'm on Windows, the GPU profiling is not yet available in WSL, can you clone my branch and run it on your side with your own benchmark in order to be sure that it looks ok.

There is a small update on the code to create the saved model:
```
import tensorflow as tf
import transformers

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("bert-base-uncased")
inputs = {"input_ids": input_ids, "attention_mask": attention_mask, "token_type_ids": token_type_ids}
base_output = base_model.bert(inputs)
seq_out, pool_out = base_output[0], base_output[1]
base_model.trainable = False
model = tf.keras.models.Model(inputs=inputs, 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")
```

No need anymore to load from the PT model, thanks to some update we have applied on the loading weights mechanism.

Hi @jplu ,

Thanks for opening the PR to integrate this change. I have cloned your branch and run benchmark on my side. The results are as expected!

**latency: (ms/batch) on 1 GPU-V100**
Batch Size | current master | tf-einsumdense |
-- | -- | --
1 | 20.9 | 6.26
2 | 24.1 | 8.68
4 | 27.6 | 13.1
8 | 36.3 | 21.5
16 | 58.8 | 42.3
32 | 94.7 | 80.4
64 | 170 | 156
128 | 321 | 309

And on GPU profiling, obtained the same observation like I posted in the very first comment: GPU computation is continuous and compact. It is not cut off by CPU process anymore.

**current master, batch_size=128**
![image](https://user-images.githubusercontent.com/70337521/102310915-c580fb80-3f20-11eb-8062-b89e8b11dc49.png)

**tf-einsumdense, batch_size=128**
![image](https://user-images.githubusercontent.com/70337521/102311089-12fd6880-3f21-11eb-9129-100081eef716.png)