Passing optimizer to Trainer constructor does not work
Author: quantitative-technologies
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.20.1
- Platform: Linux-5.4.188+-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.7.13
- Huggingface_hub version: 0.8.1
- PyTorch version (GPU?): 1.12.0+cu102 (False)
- Tensorflow version (GPU?): 2.8.2 (False)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: TPU (used on the platform, nothing specific in the script)
- Using distributed or parallel set-up in script?: No


### Who can help?

@sgugger

### Information

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

### Tasks

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

### Reproduction

Run the following script, which passes an optimizer object to `Trainer`.

```
import numpy as np
import site
import torch
from datasets import load_dataset, load_metric
from transformers import AutoModelForSequenceClassification, AutoTokenizer, Trainer, TrainingArguments, set_seed
from transformers.trainer_pt_utils import get_parameter_names

PASS_OPTIMIZER_TO_TRAINER = True

MODEL_NAME = 'albert-large-v2'
TASK = 'rte'
MAX_SEQ_LENGTH = 128
EPOCHS = 1
LEARNING_RATE = 2e-5
SEED = 10000
OPTIMIZER = 'adamw_torch'
OUTPUT_DIR = 'output'

train_args = TrainingArguments(num_train_epochs=EPOCHS, 
                               learning_rate=LEARNING_RATE,
                               seed=SEED,
                               optim=OPTIMIZER,
                               output_dir=OUTPUT_DIR,
                               overwrite_output_dir=True,
                               evaluation_strategy='epoch',
                               do_eval=True,
                               full_determinism=True)

set_seed(SEED)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
raw_datasets = load_dataset("glue", TASK)
metric = load_metric("glue", TASK)


def compute_metrics(p):
    preds = p.predictions
    preds = np.argmax(p.predictions, axis=1)
    return metric.compute(predictions=preds, references=p.label_ids)
 
def preprocess_function(examples):
    # Tokenize the texts
    args = (
        (examples['sentence1'], examples['sentence2'])
    )
    return tokenizer(*args, padding="max_length", max_length=MAX_SEQ_LENGTH, truncation=True)

raw_datasets = raw_datasets.map(
    preprocess_function,
    batched=True)


train_dataset = raw_datasets["train"]
eval_dataset = raw_datasets["validation"]

if not PASS_OPTIMIZER_TO_TRAINER:
    trainer = Trainer(
        model=model, 
        args=train_args, 
        train_dataset=train_dataset, 
        eval_dataset=eval_dataset, 
        compute_metrics=compute_metrics, 
        tokenizer=tokenizer) 

# Create adamw_torch optimizer manually
decay_parameters = get_parameter_names(model, [torch.nn.LayerNorm])
decay_parameters = [name for name in decay_parameters if "bias" not in name]
optimizer_grouped_parameters = [
    {
        "params": [p for n, p in model.named_parameters() if n in decay_parameters],
        "weight_decay": train_args.weight_decay,
    },
    {
        "params": [p for n, p in model.named_parameters() if n not in decay_parameters],
        "weight_decay": 0.0,
    },
]
optimizer = torch.optim.AdamW(optimizer_grouped_parameters,
                              lr=train_args.learning_rate,
                              betas=(train_args.adam_beta1, train_args.adam_beta2),
                              eps=train_args.adam_epsilon)

if PASS_OPTIMIZER_TO_TRAINER:
    trainer = Trainer(
        model=model, 
        args=train_args, 
        train_dataset=train_dataset, 
        eval_dataset=eval_dataset, 
        compute_metrics=compute_metrics, 
        tokenizer=tokenizer,
        optimizers=(optimizer, None))
else:
    #trainer.optimizer = optimizer
    pass

trainer.train()
```

The model fails to train. Also the training pass runs at about 2x the normal speed.

If the variable `PASS_OPTIMIZER_TO_TRAINER` is now set to `False`, the `Trainer` creates its optimizer based on `train_args`, which should be identical to the manually created one. However, now the training is successful.

I'm guessing that after passing `model` into the `Trainer` constructor it gets modified and the `optimizer` parameters are no longer valid. This is because in the script (with `PASS_OPTIMIZER_TO_TRAINER = False`) if line -4 is uncommented it has no effect, indicating that `optimizer` is now the same as `trainer.optimizer`. 

### Expected behavior

The script should work properly as is. It should have identical results whether `PASS_OPTIMIZER_TO_TRAINER` is `True` or `False`. 

If my guess is correct, then I don't see how the optimizer argument of `Trainer` can accept an optimizer object. But then that creates issues for anyone wanting to use `Trainer` with a custom optimizer. 

Comments:
cc @muellerzr, would you like to take a look at this while Sylvain is on leave?

@quantitative-technologies your issue here is how you have your if/else setup. I believe by reinstantiating an optimizer always even if we don't use it in the trainer, the model gets linked to that optimizer instead (or too?), and as a result you aren't training well. 

Here's my refactor of your code, and below shows that when doing `PASS_OPTIMIZER_TO_TRAINER` = `True` or `False`, I get the exact same results (bar the timing ever so slightly)

```python
import numpy as np
import site
import torch
from datasets import load_dataset, load_metric
from transformers import AutoModelForSequenceClassification, AutoTokenizer, Trainer, TrainingArguments, set_seed
from transformers.trainer_pt_utils import get_parameter_names

PASS_OPTIMIZER_TO_TRAINER = True

MODEL_NAME = 'albert-large-v2'
TASK = 'rte'
MAX_SEQ_LENGTH = 128
EPOCHS = 1
LEARNING_RATE = 2e-5
SEED = 10000
OPTIMIZER = 'adamw_torch'
OUTPUT_DIR = 'output'


train_args = TrainingArguments(
    num_train_epochs=EPOCHS, 
    learning_rate=LEARNING_RATE,
    seed=SEED,
    optim=OPTIMIZER,
    output_dir=OUTPUT_DIR,
    overwrite_output_dir=True,
    evaluation_strategy='epoch',
    do_eval=True,
    full_determinism=True
)

set_seed(SEED)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
raw_datasets = load_dataset("glue", TASK)
metric = load_metric("glue", TASK)


def compute_metrics(p):
    preds = p.predictions
    preds = np.argmax(p.predictions, axis=1)
    return metric.compute(predictions=preds, references=p.label_ids)
 
def preprocess_function(examples):
    # Tokenize the texts
    args = (
        (examples['sentence1'], examples['sentence2'])
    )
    return tokenizer(*args, padding="max_length", max_length=MAX_SEQ_LENGTH, truncation=True)

raw_datasets = raw_datasets.map(
    preprocess_function,
    batched=True
)


train_dataset = raw_datasets["train"]
eval_dataset = raw_datasets["validation"]

if not PASS_OPTIMIZER_TO_TRAINER:
    trainer = Trainer(
        model=model, 
        args=train_args, 
        train_dataset=train_dataset, 
        eval_dataset=eval_dataset, 
        compute_metrics=compute_metrics, 
        tokenizer=tokenizer
)

else:
    # Create adamw_torch optimizer manually
    decay_parameters = get_parameter_names(model, [torch.nn.LayerNorm])
    decay_parameters = [name for name in decay_parameters if "bias" not in name]
    optimizer_grouped_parameters = [
        {
            "params": [p for n, p in model.named_parameters() if n in decay_parameters],
            "weight_decay": train_args.weight_decay,
        },
        {
            "params": [p for n, p in model.named_parameters() if n not in decay_parameters],
            "weight_decay": 0.0,
        },
    ]
    optimizer = torch.optim.AdamW(
        optimizer_grouped_parameters,
        lr=train_args.learning_rate,
        betas=(train_args.adam_beta1, train_args.adam_beta2),
        eps=train_args.adam_epsilon
    )

    trainer = Trainer(
        model=model, 
        args=train_args, 
        train_dataset=train_dataset, 
        eval_dataset=eval_dataset, 
        compute_metrics=compute_metrics, 
        tokenizer=tokenizer,
        optimizers=(optimizer, None)
    )

trainer.train()
```

With passing: 
```python
{'epoch': 1.0,
 'train_loss': 0.5830265436417017,
 'train_runtime': 206.1904,
 'train_samples_per_second': 12.076,
 'train_steps_per_second': 1.513}
```
Without passing:
```python
{'epoch': 1.0,
 'train_loss': 0.5830265436417017,
 'train_runtime': 205.4122,
 'train_samples_per_second': 12.122,
 'train_steps_per_second': 1.519}
```

@muellerzr I don't believe that was the issue. My if/else structure was convoluted because I was trying to make an additional point besides the bug. 

Now when I run the exact code you sent, I am still seeing the original problem:

```
{'train_runtime': 50.4719, 
 'train_samples_per_second': 49.334, 
 'train_steps_per_second': 6.182, 
 'train_loss': 0.7272015596047426, 
 'epoch': 1.0}
```

So it is not learning. Also mine is running ~4x faster than yours, though we may be using different hardware.

I am using a Colab instance with a TPU (using a single-core, i.e. not distributed for the purpose of the bug report). 

Could the TPU be an issue here? Can you test it out on a TPU?

@quantitative-technologies I could recreate your issue once I had torch-xla installed. Looking into it

@quantitative-technologies the issue is described in this xla issue: https://github.com/pytorch/xla/issues/3675#issuecomment-1171702988

Essentially you need to place the model on the device yourself first, then create the optimizer. After doing this I got the exact same results. 

*Super* subtle bug here. 

To do so, add the following lines to your code:

```python
import torch_xla.core.xla_model as xm

# In the if/else of to create the optimizer yourself
    # Create adamw_torch optimizer manually
    model = model.to(xm.xla_device()) # <- We need to move the model to the device *before* creating the optimizer
    decay_parameters = get_parameter_names(model, [torch.nn.LayerNorm])
    decay_parameters = [name for name in decay_parameters if "bias" not in name]
    ...
```


> @quantitative-technologies I could recreate your issue once I had torch-xla installed. Looking into it

Right, I should have mentioned that `torch-xla` was installed in the report.

> @quantitative-technologies the issue is described in this xla issue: [pytorch/xla#3675 (comment)](https://github.com/pytorch/xla/issues/3675#issuecomment-1171702988)
> 
> Essentially you need to place the model on the device yourself first, then create the optimizer. After doing this I got the exact same results.

OK, I see. 

I was impatient and made my own solution, by subclassing `Trainer` with an `optimizers_init` function argument that creates the optimizer from the `Trainer`'s model and the `args`. This resolved the issue, since `Trainer` places the model for you.

I'm not sure if this is perhaps an improved design. Actually, having an `lr_scheduler_init` definitely makes more sense -- I didn't implement it yet though. This is because the `Trainer` does the calculation for the number of training steps which is needed to build the `lr_regularizer` anyhow. 

If there is interest updating the `Trainer` I can submit a push request. Here is my code for my subclass:

```
class TrainerOptimizerInit(Trainer):
    """
    Args:
        optimizers_init (`Tuple[Callable[[Union[PreTrainedModel, nn.Module], TrainingArguments], torch.optim.Optimizer], 
                                torch.optim.lr_scheduler.LambdaLR]`, *optional*): A tuple containing (1) a function that is
            used to create an optimizer from the `model` and `args`, and (2) the scheduler to use. Will default to an 
            instance of [`AdamW`] on your model and a scheduler given by [`get_linear_schedule_with_warmup`] controlled 
            by `args`.
    """
    def __init__(
        self,
        model: Union[PreTrainedModel, nn.Module] = None,
        args: TrainingArguments = None,
        data_collator: Optional[DataCollator] = None,
        train_dataset: Optional[Dataset] = None,
        eval_dataset: Optional[Dataset] = None,
        tokenizer: Optional[PreTrainedTokenizerBase] = None,
        model_init: Callable[[], PreTrainedModel] = None,
        compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
        callbacks: Optional[List[TrainerCallback]] = None,
        optimizers_init: Tuple[Callable[[Union[PreTrainedModel, nn.Module], TrainingArguments], torch.optim.Optimizer], 
                               torch.optim.lr_scheduler.LambdaLR] = (None, None),
        preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = None,
    ):
        super().__init__(model=model, 
                         args=args, 
                         data_collator=data_collator, 
                         train_dataset=train_dataset, 
                         eval_dataset=eval_dataset, 
                         tokenizer=tokenizer, 
                         model_init=model_init, 
                         compute_metrics=compute_metrics, 
                         callbacks=callbacks,
                         preprocess_logits_for_metrics=preprocess_logits_for_metrics)

        self.optimizer_init, self.lr_scheduler = optimizers_init

    def create_optimizer(self):
        """
        Setup the optimizer.

        We provide a reasonable default that works well. If you want to use something else, you can subclass and override 
        this method in a subclass.
        """
        opt_model = self.model_wrapped if is_sagemaker_mp_enabled() else self.model

        if self.optimizer is None:
            if self.optimizer_init is None:
                # fall back to original behaviour
                return super().create_optimizer()

            self.optimizer = self.optimizer_init(opt_model, self.args)

        if is_sagemaker_mp_enabled():
            self.optimizer = smp.DistributedOptimizer(self.optimizer)

        return self.optimizer
```


That'd be up to @sgugger, who is OOF until the 30th on holiday 😄 I'll make sure he sees this though when he's back!

We don't really want to add other init functions. For any specific behavior, users should subclass the `create_optimizer` function directly.

I am confused. Should the optimizer be passed to TrainingArguments or to the Trainer?

@sgugger @muellerzr 
is there any documentation available around it, I have requirement to use  Torch optim scheduler such as cosine with restarts i am unable to find a way to pass it, we have been given gette fucntions but how to pass to the trainer..
any eg. will be help ful.


You can pass the `optimiser` as an argument to the trainer. See the documentation [here](https://huggingface.co/docs/transformers/main_classes/trainer#transformers.Trainer.optimizers) as well as the following: 
```python 
    def create_optimizer_and_scheduler(self, num_training_steps: int):
        """
        Setup the optimizer and the learning rate scheduler.

        We provide a reasonable default that works well. If you want to use something else, you can pass a tuple in the
        Trainer's init through `optimizers`, or subclass and override this method (or `create_optimizer` and/or
        `create_scheduler`) in a subclass.
        """
        self.create_optimizer()
        if IS_SAGEMAKER_MP_POST_1_10 and smp.state.cfg.fp16:
            # If smp >= 1.10 and fp16 is enabled, we unwrap the optimizer
            optimizer = self.optimizer.optimizer
        else:
            optimizer = self.optimizer
        self.create_scheduler(num_training_steps=num_training_steps, optimizer=optimizer)
```
 Thus subclassing  `create_optimizer` will help you create  which ever optimizer. Does that make sense for you?