Adam Weight Decay Rate does not hear the opinion of tf.stop_gradient
Author: arivero
Pull request: no

### System Info

- `transformers` version: 4.26.1
- Platform: Linux-5.19.0-32-generic-x86_64-with-glibc2.35
- Python version: 3.10.6
- Huggingface_hub version: 0.12.1
- PyTorch version (GPU?): 1.11.0+cu113 (True)
- Tensorflow version (GPU?): 2.9.0 (True)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: YES
- Using distributed or parallel set-up in script?: NO

### Who can help?

@gante @Rocketknight1 @sgugger 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

I use a gpt2 variant. 
`model = TFGPT2LMHeadModel.from_pretrained("PlanTL-GOB-ES/gpt2-large-bne", from_pt=True)`
where I freeze all the transformer weights except wte
`for i in range(36):  model.transformer.h[i].trainable=False`
and then monkey patch the wte layer to stop propagation of gradiente for only some tokens. For instance:
```
import transformers
from transformers.tf_utils import shape_list
def call(self, inputs: tf.Tensor, mode: str = "embedding") -> tf.Tensor:
    w= tf.stop_gradient(self.maskEmbeddings * self.weight) + (1-self.maskEmbeddings) * self.weight
    if mode == "embedding":
            return tf.gather(w, inputs)
    elif mode == "linear":
            first_dims = shape_list(inputs)[:-1]
            x = tf.reshape(inputs, [-1, self.hidden_size])
            logits = tf.matmul(x, w, transpose_b=True)
            return tf.reshape(logits, first_dims + [self.vocab_size])
    else:
            raise ValueError(f"mode {mode} is not valid.")
transformers.modeling_tf_utils.TFSharedEmbeddings.call = call
```
key line here is `tf.stop_gradient(self.maskEmbeddings * self.weight`

Now I provide some mask matrix model.transformer.wte.maskEmbeddings = maskEmbeddings and run the training as usual with AdamWeightDecay optimizer. I check the changes with a checksum:
```
def printChecksum(model):
    frozenWeights=model.transformer.wte.maskEmbeddings*model.transformer.wte.weight
    check=tf.reduce_sum(frozenWeights, axis=0)
    print(check)
printChecksum(model)
```
when weight_decay_rate is different of 0.0, the optimizer applies decay to the weights. 


### Expected behavior

I would expect the weights that have opted out of gradient updating via the tf.stop_gradient to remain unaltered even if weight_decay_rate is not zero.

On other hand, it is true that the wdr is not part of the gradient, so it should be stopped by other ways. Surely the right path is to implement support for soft prompt training at a higher level.

Comments:
This is a standard behaviour of all TF optimizers - `tf.stop_gradient()` stops gradient "flowing" back through that operation, but the optimizer is not aware of this control step, and so it simply sees weights with a gradient of 0 and applies weight decay to them as normal (in fact, it probably also applies updates to them from the residual momentum in the optimizer).

If you want to exclude weights from being updated entirely, you can set `layer.trainable=False`, which you're already doing. However, it sounds like you want to selectively mask different weights in each step, and you definitely can't set properties like `layer.trainable` in the `call()` function.

In those circumstances, probably the best solution is to override `train_step()` instead and apply the optimizer step to only unmasked weights at each stage (note that this might prevent XLA compilation!) Alternatively, you could just use a standard optimizer that doesn't have weight decay.

Soft prompt tuning is definitely something we could investigate adding as a feature, though! What kind of interface do you think would work for it?

I see. Was not sure if it was a bug, just as you say, one needs to know what the standard behaviour is. 

Should I open a feature request for soft prompt? It would definitely be an interesting addition to the toolset, and there is a lot of tripwires that one can trigger if made in an ad-hoc way. This was not the only one :-)

cc @gante to that one - I believe he's been looking at soft prompting in generation!

Hey @arivero 👋 If I got it right, you would be interested in soft-prompting as in passing post-embedding values to the model (and not train on a few masked tokens). 

Our models support an `input_embeds` input, which is mutually exclusive with `input_ids` and is meant as post-embeddings `input_ids`. Would using this input instead solve your problem? I believe you would be able to handle all masking operations outside the model :)