Skip to content

TF LMHead very slow: TFGPT2LMHeadModel is 7 times slower than Torch GPT2LMHeadModel  #6264

Description

@gyin94

The benchmark shows the Tensorflow TFGPT2LMHeadModel is 7 times (140ms) slower than Torch GPT2LMHeadModel implementation (20ms). Based on the profile tool from latest https://www.tensorflow.org/tfx/serving/tensorboard, the bottleneck is here.

logits = tf.matmul(x, self.weight, transpose_b=True)

from transformers import *

args = PyTorchBenchmarkArguments(
    models=["gpt2_pt"],
    batch_sizes=[1], sequence_lengths=[8])

config_base = GPT2Config(n_layer=2, n_head=2)
config_base.architectures = ["GPT2LMHeadModel"]
benchmark = PyTorchBenchmark(args, configs=[config_base])
print(benchmark.run())

args = TensorFlowBenchmarkArguments(
    models=["gpt2_tf"],
    batch_sizes=[1], sequence_lengths=[8])

config_base = GPT2Config(n_layer=2, n_head=2)
config_base.architectures = ["GPT2LMHeadModel"]
benchmark = TensorFlowBenchmark(args, configs=[config_base])
print(benchmark.run())

Without LMHeadModel, the benchmark is 6ms vs 8ms.

Metadata

Metadata

Labels

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions