apply_gradients() return error using official.nlp.optimization

[sehamra]

System information
OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Google Colab
TensorFlow version and how it was installed (source or binary): tf-nightly 2.2.0-dev20200408 (pip install) and 2.1 stable version
TensorFlow-Addons version and how it was installed (source or binary):
Python version:
Is GPU used? (yes/no): TPU
Describe the bug

A clear and concise description of what the bug is.

Code to reproduce the issue

optimizer = optimization.create_optimizer(
    init_lr=INIT_LR,
    num_train_steps=NB_BATCHES_TRAIN, # per epochs
    num_warmup_steps=WARMUP_STEPS)

def squad_loss_fn(labels, model_outputs):
    start_positions = labels['start_positions']
    end_positions = labels['end_positions']
    start_logits, end_logits = model_outputs

    start_loss = tf.keras.backend.sparse_categorical_crossentropy(
        start_positions, start_logits, from_logits=True)
    
    end_loss = tf.keras.backend.sparse_categorical_crossentropy(
        end_positions, end_logits, from_logits=True)
    
    total_loss = (tf.reduce_mean(start_loss) + tf.reduce_mean(end_loss)) / 2
    return total_loss

train_loss = tf.keras.metrics.Mean(name="train_loss")

bert_squad.compile(optimizer,
                   squad_loss_fn)

# Training loop
NB_EPOCHS = 3
for epoch in range(NB_EPOCHS):
    print("Start of epoch {}".format(epoch+1))
    start = time.time() 
    
    train_loss.reset_states() 
    
    for (batch, (inputs, targets)) in enumerate(train_dataset_light):
        with tf.GradientTape() as tape:
          model_outputs = bert_squad(inputs)
          loss = squad_loss_fn(targets, model_outputs)
        grads = tape.gradient(loss, bert_squad.trainable_variables) 
        optimizer.apply_gradients(zip(grads, bert_squad.trainable_variables))
        #optimizer.apply_gradients(zip(grads, bert_squad.trainable_variables), name=None, all_reduce_sum_gradients=True))
        
        train_loss(loss)

Other info / logs

same issue here : Missing argument in apply_gradients() in AdamW optimizer #1267

[gabrieldemarmiesse]

Thanks for the report, what is the version of tensorflow-addons that you are using?
You can do tfa.__version__ or !pip freeze for that.

[prasad4fun]

Facing the same issue in tf-nightly 2.2.0-dev20200411

OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Google Colab
Python version:3.6.9
tensorflow-addons==0.8.3
Is GPU used? (yes/no): TPU

[gabrieldemarmiesse]

Could you try with tensorflow-addons 0.9.x? (the latest version).

[prasad4fun]

Same issue even after upgrading to tensorflow-addons 0.9.0

Signature: optimizer.apply_gradients(self,grads_and_vars,name=None,all_reduce_sum_gradients=True):

def apply_gradients(self,
                      grads_and_vars,
                      name=None,
                      all_reduce_sum_gradients=True):
    grads, tvars = list(zip(*grads_and_vars))
    (grads, _) = tf.clip_by_global_norm(grads, clip_norm=1.0)
    return super(AdamWeightDecay, self).apply_gradients(
        zip(grads, tvars),
        name=name,
        all_reduce_sum_gradients=all_reduce_sum_gradients)

Signature: tf.keras.optimizers.Adam.apply_gradients(grads_and_vars, name=None, experimental_aggregate_gradients=True)

def apply_gradients(self,
                      grads_and_vars,
                      name=None,
                      experimental_aggregate_gradients=True):

keyword argument 'all_reduce_sum_gradients' likely replaced with 'experimental_aggregate_gradients'

[keshan]

Same issue even after upgrading to tensorflow-addons 0.9.0

Signature: optimizer.apply_gradients(self,grads_and_vars,name=None,all_reduce_sum_gradients=True):

def apply_gradients(self,
                      grads_and_vars,
                      name=None,
                      all_reduce_sum_gradients=True):
    grads, tvars = list(zip(*grads_and_vars))
    (grads, _) = tf.clip_by_global_norm(grads, clip_norm=1.0)
    return super(AdamWeightDecay, self).apply_gradients(
        zip(grads, tvars),
        name=name,
        all_reduce_sum_gradients=all_reduce_sum_gradients)

Signature: tf.keras.optimizers.Adam.apply_gradients(grads_and_vars, name=None, experimental_aggregate_gradients=True)

def apply_gradients(self,
                      grads_and_vars,
                      name=None,
                      experimental_aggregate_gradients=True):

keyword argument 'all_reduce_sum_gradients' likely replaced with 'experimental_aggregate_gradients'

You are right. 'all_reduce_sum_gradients' have been renamed to 'experimental_aggregate_gradients'
tensorflow/tensorflow#37908

[seanpmorgan]

So I believe this is an issue with the version of tensorflow/models you're using. The official BERT nlp implementation has an option of two optimizers (AdamWeightDecay and LAMB).
https://github.com/tensorflow/models/blob/master/official/nlp/optimization.py#L87-L96

LAMB is the optimizer in TFA:
https://github.com/tensorflow/addons/blob/master/tensorflow_addons/optimizers/lamb.py#L30

But AdamWeightDecay is actually an optimizer in the NLP repo itself:
https://github.com/tensorflow/models/blob/d7a44150b83cc5abbfcee4c16f6f9e0ee642379d/official/nlp/optimization.py#L111

Per the model team AdamWeightDecay optimizer is not the same as TFA's AdamW since there are some configurations that are particular to that repository. We apologize for the duplicated functionality and are actively working with the model team to clean up as much as we can. However, it looks like a fix for the issue you're describing has already landed in the tensorflow/models repo so I would suggest updating your clone:
tensorflow/models@9fc4fd0

As far as this issue goes for TFA. LAMB optimizer simply subclasses tf.keras.optimizers.Optimizer and makes no modifications to _apply_gradients so it should be stable to this change. Our AdamW optimizer supports this new experimental argument through the kwargs addition to our call:
#1566

Closing, but feel free to re-open if you don't feel this fixed your issue.

CC @saberkun just as a data point to show how having this functionality in two repositories is causing a poor user experience