[TF Longformer] Improve Speed for TF Longformer #6447
Merged
Conversation
This file contains 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
Codecov Report
@@ Coverage Diff @@
## master #6447 +/- ##
==========================================
+ Coverage 78.96% 79.81% +0.84%
==========================================
Files 157 157
Lines 28486 28479 -7
==========================================
+ Hits 22495 22730 +235
+ Misses 5991 5749 -242
Continue to review full report at Codecov.
|
patrickvonplaten
force-pushed
the
add_tf_compile_function_to_test
branch
from
8c8c5a0 to
85bed01
Compare
| # is index masked or global attention | ||
|
|
||
| is_index_masked = tf.math.less(attention_mask, 0) | ||
| is_index_global_attn = tf.math.greater(attention_mask, 0) |
There was a problem hiding this comment.
These operations should be done once at not be repeated for every layer -> unnecessary slow-down
| @@ -170,24 +168,20 @@ def call( | |||
| # normalize query | |||
| query_vectors /= tf.math.sqrt(tf.constant(self.head_dim, dtype=tf.dtypes.float32)) | |||
|
|
|||
| query_vectors = tf.transpose( | |||
There was a problem hiding this comment.
get rid of tf.transpose as it allocates a new tensor upon calling. Should be avoided.
| @@ -134,6 +125,19 @@ def test_save_load(self): | |||
|
|
|||
| self.assert_outputs_same(after_outputs, outputs) | |||
|
|
|||
| def test_graph_mode(self): | |||
There was a problem hiding this comment.
Add a function that tests every model in graph mode
patrickvonplaten
changed the title
Speed Benchmarking:Running this command on the master branch: on this env: gives: On this branch the speed is improved to: So we can see an improvement of ca. 3%, which is not that much actually... I guess it's interesting to see what effect removing some unnecessary |
patrickvonplaten
requested review from
jplu,
sgugger and
LysandreJik
and removed request for
jplu
patrickvonplaten
changed the title
sgugger
approved these changes
Aug 14, 2020
patrickvonplaten
force-pushed
the
add_tf_compile_function_to_test
branch
from
316206c to
ae3bbe2
Compare
LysandreJik
approved these changes
Aug 26, 2020
Zigur
pushed a commit
to Zigur/transformers
that referenced
this pull request
Oct 26, 2020
* add tf graph compile tests * fix conflict * remove more tf transpose statements * fix conflicts * fix comment typos * move function to class function * fix black * fix black * make style
fabiocapsouza
pushed a commit
to fabiocapsouza/transformers
that referenced
this pull request
Nov 15, 2020
* add tf graph compile tests * fix conflict * remove more tf transpose statements * fix conflicts * fix comment typos * move function to class function * fix black * fix black * make style
fabiocapsouza
added a commit
to fabiocapsouza/transformers
that referenced
this pull request
Nov 15, 2020
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
This PR:
tf.transpose()(In contrast to PyTorch,tf.transpose()allocates a new tensor and thus should be avoided). This also cleans up the code IMO.=> These changes lead to a speed-up of 1.03 which is actually not that much...more details in benchmark below.
After a lot of digging TF XLA will not be very easy to implement as a lot of kernels that are highly used in this model
tf.whereare not implemented for XLA (yet). So TF Longformer TPU will not work sadly for the moment @ibeltagyConclusion
For me the PR was also a good exercise to see whether TF can significantly sped up by removing unnecessary tensor allocations. It seems like it's not really worth it go through all the tf models if the improvement in speed is only around 2,3%.