tensorflow: Tensorflow 2.3.0 is much slower than PyTorch 1.4.0 in backward propagation. (20 times slower)

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): Yes
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Win10
• CUDA/cuDNN version: CUDA 10.1/cuDNN 7
• TensorFlow version (use command below): 2.3.0
• Python version: 3.7.4
• GPU model and memory: Geforce GTX 1660Ti (6GB)

Describe the current behavior

• I tried to use both TF2 and PyTorch to train SkipGram model with negative sampling. But notice a huge training speed difference.

• TF2 takes ~0.5s to finish one step while PyTorch takes ~0.01s. One step refers to use 1 batch data to do forward and backward and update the weights.

• The main difference is in backward propagation.

• More specifically, ‘apply_gradients()’ in TF2 is much slower than counterpart ‘step()’ in PyTorch. Does anyone know why?

• Notes:

1. Both versions use GPU.
2. tf.function is used.
3. Training batch size and data is exactly the same.
4. Model size (vocab_size, emb_size) is the same.

• TF2:

@tf.function
def train_step_tf(pos_w, pos_v, neg_v):
    print('-'*50)
    with tf.GradientTape() as tape:  
        loss = skip_gram_model_tf(pos_w, pos_v, neg_v)
    
    start_time = time.time()
    variables = skip_gram_model_tf.trainable_variables
    print('getting variables time = {}'.format(time.time() - start_time))

    start_time = time.time()
    gradients = tape.gradient(loss, variables)
    print('init gradient time = {}'.format(time.time() - start_time))

    start_time = time.time()
    optimizer.apply_gradients(zip(gradients, variables))  # !!!!! this is where the major speed difference happens!!!!!
    print('apply gradient time = {}'.format(time.time() - start_time))
    
    print('-'*50)

    return loss

• PyTorch:

def train_step_torch(pos_w, pos_v, neg_v):  
    print('-'*50)
    start_time = time.time()
    optimizer_torch.zero_grad()
    loss = skip_gram_model_torch.forward(pos_w, pos_v, neg_v)
    print('getting variables time = {}'.format(time.time() - start_time))

    start_time = time.time()
    loss.backward()
    print('init gradient time = {}'.format(time.time() - start_time))
    
    start_time = time.time()
    optimizer_torch.step() # !!!!this is fast!!!!
    print('apply gradient time = {}'.format(time.time() - start_time))
    print('-'*50)

    return loss

Describe the expected behavior I expect the performance shouldn’t be that different. Is there anything I miss to use in TF2?

Standalone code to reproduce the issue Colab link of full code: https://colab.research.google.com/drive/17QsTkV271LPvo6aJf1QOGuumdU8OXmVj?usp=sharing