Deepspeed integration ignores Optuna trial parameters in hyperparameter_search
Author: Mindful
Pull request: no
Labels: name: DeepSpeed

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.6.1
- Platform: Linux-4.19.0-16-cloud-amd64-x86_64-with-glibc2.10
- Python version: 3.8.10
- PyTorch version (GPU?): 1.8.1+cu111 (True)
- Tensorflow version (GPU?): not installed (NA)
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: no

Also in case it matters, my deepspeed version is 0.3.16.

### Who can help

Maybe @stas00 ? I'm not sure.

## Information

Model I am using: custom Pytorch transformer model, although I don't think it matters here.

The problem arises when using:
* [ ] the official example scripts: (probably, I haven't tried)
* [x] my own modified scripts: I'm running a simple custom MLM training script using the transformers trainer.

## To reproduce

Steps to reproduce the behavior:

1. Run a hyperparameter search using the transformers Trainer with the [default zero-2 config from the documentation](https://huggingface.co/transformers/main_classes/trainer.html#zero-2-example)
2. Observe that parameters taken from the Deepspeed config like gradient accumulation steps and optimizer/scheduler params are not updated to reflect the Optuna trial parameters.

Here is some output from the script I'm running, with the middle omitted for brevity. I'm printing trial params myself but I do it from inside the Trainer, so these are definitely the same trial params the Trainer is getting.

```
[I 2021-05-27 02:20:41,133] A new study created in memory with name: no-name-26248906-22c0-4666-a7d4-159173902bc5
current trial params {'learning_rate': 2.0670100636747183e-05, 'adam_beta2': 0.98, 'gradient_accumulation_steps': 6, 'dropout': 0.0, 'local_attention_window': 192, 'weight_decay': 0.01, 'warmup_ratio': 0.12, 'deep_transformer_stack_layers': 10}

.....

[2021-05-27 02:21:01,787] [INFO] [config.py:751:print]   initial_dynamic_scale ........ 65536
[2021-05-27 02:21:01,787] [INFO] [config.py:751:print]   loss_scale ................... 0
[2021-05-27 02:21:01,787] [INFO] [config.py:751:print]   memory_breakdown ............. False
[2021-05-27 02:21:01,787] [INFO] [config.py:751:print]   optimizer_legacy_fusion ...... False
[2021-05-27 02:21:01,787] [INFO] [config.py:751:print]   optimizer_name ............... adamw
[2021-05-27 02:21:01,787] [INFO] [config.py:751:print]   optimizer_params ............. {'lr': 5e-05, 'betas': [0.9, 0.999], 'eps': 1e-08, 'weight_decay': 0.0}
[2021-05-27 02:21:01,787] [INFO] [config.py:751:print]   pipeline ..................... {'stages': 'auto', 'partition': 'best', 'seed_layers': False, 'activation_checkpoint_interval': 0}
[2021-05-27 02:21:01,788] [INFO] [config.py:751:print]   pld_enabled .................. False

....
```
I looked around briefly and I _think_ the issue comes from the fact that that the Deepspeed config [is built as part of TrainingArgs](https://github.com/huggingface/transformers/blob/996a315e76f6c972c854990e6114226a91bc0a90/src/transformers/training_args.py#L677) and then presumably never updated after that, even if the training args change. Consequently when the [training args are updated](https://github.com/huggingface/transformers/blob/996a315e76f6c972c854990e6114226a91bc0a90/src/transformers/trainer.py#L861) as part of setup for the hyperparameter search, it's not reflected in the Deepspeed config.

Note that this might also be an issue with Ray, I just haven't tried it with Ray. 

## Expected behavior

Ideally Deepspeed would run with config/parameters that respected the content of the Optuna trials, although I know that getting two external integrations to play well together is easier said than done. In the meantime I'm going to see if I can work around this by using an HF scheduler and HF optimizer in the hopes that those will take their parameters from the training arguments directly.



Comments:
Thanks for the report, @Mindful!

It's very possible that I mistakenly bypassed some logic of optuna, as I have never used it. 

Would you like to have a look and see if you can fix it - since you have already been researching this - basically go into `src/transformers/trainer.py` look for `optuna` and see where `if self.deepspeed` skips it. Shouldn't be too difficult if you already have a test for it.

Please let me know if I can help.

Thank you!

@stas00 

I would need to look a little closer to be sure but it really just looks like a timing issue - the Deepspeed config is built off the Trainer config, which then has its state changed *afterwards* by the Optuna integration so the two get out of sync. 

I am definitely open to trying to fix this myself (I've been looking for a reason to contribute), I just have two concerns:
1. I'm pretty swamped right now, at least for the next week or two. Things should hopefully calm down after that, but it's hard for me to promise I can get to it by a certain date. 
2. It seems like the only options for fixing this are either somehow making the Deepspeed config automatically update to reflect updates to the training config (which would be complicated and probably overkill) or changing the hyperparameter_search method so that it also updates the Deepspeed config if necessary. I think the latter is the better option, but going attribute-by-attribute would basically mean duplicating the logic for copying training parameters from the TrainingArguments to the deepspeed config. I think the _best_ option is to just construct a new DeepSpeedConfigHF based on the updated training parameters, but there's a lot of logic there and I'm not sure if this is safe to do. 

Actually, if the fix is as easy as rebuilding DeepSpeedConfigHF with the updated TrainingArguments, this might be relatively quick. I'm not sure. 
Edit: I wrote this out and then went back and looked at the code and I think doing the above might fix it (in which case this is an easy fix). Let me try this and get back to you.

Yeah, I just changed files locally so that the hyperparameter search rebuilds the DeepSpeedConfigHF object and that seems to have fixed it. I still need to double check tests pass/etc, but it looks like this was much easier than I thought.

I'll open a PR shortly. 

Awesome. Looking forward to reading your PR, @Mindful - please tag me there.

For tests just run:

```
RUN_SLOW=1 pytest tests/deepspeed
```

and if you have an idea for a new test that would be a bonus.