Use DS callable API to allow hf_scheduler + ds_optimizer

[tjruwase]

This PR:

• Used the (new) Callable api of deepspeed.initialize() to enable combining hf schedulers with deepspeed optimizers.
• create_scheduler now has an optional optimizer arg
• Updates relevant unit test.

Blocking events: All unblocked now.

• depends on deepspeed PR 1316.
• needs new deepspeed version after PR is merged and need to update the dependencies when that happens.

deepspeed: @stas00.

[stas00]

Thank you for solving this limitation, @tjruwase

I just proposed 2 fixes to keep things back-compatible. Please let me know if it's still OK.

[stas00]

We will have to wait for @sgugger to finalize this as we are messing with the API. I think he's back on the 30th. in a week.

[stas00]

Thank you for making the changes, Tunji. Looking good.

Now let's wait for the merge of your DS PR and the new release.

[jeffra]

• Support Callable type for client optimizer and lr_scheduler deepspeedai/DeepSpeed#1316 is merged
• v0.5.1 released to PyPI: https://pypi.org/project/deepspeed/0.5.1/

[sgugger]

Thanks for your PR!

I have no problem with optimizer becoming an optional argument of Trainer.create_scheduler, I am less convinced by the need to change the return type of the two methods Trainer.create_optimizer and Trainer.create_scheduler, though I don't have particularly strong feelings. I understand it's more convenient to do

return trainer.create_scheduler(xxx)

instead of

trainer.create_scheduler(xxx)
return trainer.lr_scheduler

but it's internal code.

It's just a thought, if you really prefer the change of return type, I'm fine with it.