

TITLE
[DeepSpeed in notebooks] Jupyter + Colab

COMMENTS
0

REACTIONS
+1: 0
-1: 0
laugh: 0
hooray: 0
heart: 0
rocket: 0
eyes: 0

BODY
This PR addresses issues raised in https://github.com/huggingface/transformers/issues/10011 when the user tried to use DeepSpeed in a notebook (most likely colab).

This PR:

* forces device and distributed setup init from TrainingArguments explicitly at the beginning of Trainer's `__init__`. This is needed since until now the init was happening as a side effect of someone calling `device` or `n_gpus`, which doesn't happen if someone runs their own version of `Trainer` w/ deepspeed - which is the case with notebooks - so we are missing out on DeepSpeed init and things weren't working. Let's do it explicitly, and not as a side-effect, so everything is loud and clear.
* sets up `self.local_rank` based on LOCAL_RANK env var under deepspeed to save users a hassle - deepspeed must have `local_rank>-1`. I guess the fake launcher env setup could be folded into the init as well, but then the user loses control over the port number and they may need to edit it, so for now leaving it outside - but will ask deepspeed to provide a wrapper function to make it easy for the user. perhaps once the wrapper is available it could be automated completely. Alternatively, if they make `mpi4py` a dependency then fake launcher env setup won't be needed at all.
* documents how to run DeepSpeed in the notebook env
* adds a test that mocks a notebook environment and runs deepspeed w/o a launcher

I may wait to hear about the follow up to https://github.com/microsoft/DeepSpeed/issues/748 to merge this, if it looks quick, but otherwise I will revise the setup doc in a future PR. The main changes of this PR besides the doc are required anyway.

@sgugger 